Replace default BernoulliStepScheduler with TimeScheduler (defaulting to 10 minutes) for deterministic flushing.

envlogger/backends/riegeli_backend_writer.py

@@ -15,6 +15,7 @@
 
 """For writing trajectory data to riegeli files."""
 
+import datetime
 from typing import Any, Optional
 
 from absl import logging
@@ -53,14 +54,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 time-based scheduler with a 10-minute
+        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(
+          datetime.timedelta(minutes=10)
+      )
     else:
       self._flush_scheduler = flush_scheduler
     self._data_writer = riegeli_dataset_writer.RiegeliDatasetWriter()

envlogger/backends/riegeli_backend_writer_test.py

@@ -21,6 +21,7 @@
 
 import operator
 from typing import Any, Optional
+from unittest import mock
 
 from absl import logging
 from absl.testing import absltest
@@ -205,6 +206,29 @@ def test_episode_step_out_of_bounds_indices(self):
                           -len(episode) - 1)
 
 
+  def test_default_scheduler_is_time_based(self):
+    writer = riegeli_backend_writer.RiegeliBackendWriter(
+        data_directory=self.create_tempdir().full_path
+    )
+    scheduler = writer._flush_scheduler
+    self.assertIsInstance(scheduler, schedulers.TimeScheduler)
+    self.assertEqual(scheduler._interval_seconds, 600.0)
+
+  def test_flush_scheduler_is_called(self):
+    mock_flush_scheduler = mock.create_autospec(
+        schedulers.TimeScheduler, instance=True, return_value=True
+    )
+
+    writer = riegeli_backend_writer.RiegeliBackendWriter(
+        data_directory=self.create_tempdir().full_path,
+        flush_scheduler=mock_flush_scheduler,
+    )
+
+    step = step_data.StepData(dm_env.restart(observation=None), action=None)
+    writer.record_step(step, is_new_episode=True)
+
+    mock_flush_scheduler.assert_called_once_with(step)
+
 
 if __name__ == '__main__':
   absltest.main()

envlogger/backends/schedulers.py

@@ -16,6 +16,8 @@
 """Common logging scheduling strategies."""
 
 from collections.abc import Callable
+import datetime
+import time
 
 from envlogger import step_data
 import numpy as np
@@ -155,3 +157,33 @@ def __call__(self, data: step_data.StepData):
     if data.timestep.last():
       self._episode_counter += 1
     return should_log
+
+
+class TimeScheduler:
+  """Returns `True` after a specific time interval has elapsed."""
+
+  def __init__(
+      self,
+      interval: datetime.timedelta | float | int,
+      clock: Callable[[], float] = time.monotonic,
+  ):
+    """Initializes the time scheduler class."""
+
+    if isinstance(interval, datetime.timedelta):
+      self._interval_seconds = interval.total_seconds()
+    else:
+      self._interval_seconds = float(interval)
+
+    if self._interval_seconds <= 0:
+      raise ValueError(f'interval must be positive, got {interval!r}')
+
+    self._clock = clock
+    self._last_trigger_time = self._clock()
+
+  def __call__(self, unused_data: step_data.StepData | None = None) -> bool:
+    """Returns `True` if the interval has elapsed since the last trigger."""
+    now = self._clock()
+    if now - self._last_trigger_time >= self._interval_seconds:
+      self._last_trigger_time = now
+      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."""
+import datetime
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -363,6 +363,163 @@ def test_list_episodes_logspace(self):
         for timestep in _create_episode(num_transitions=np.random.randint(100)):
           self.assertFalse(scheduler(timestep))
 
+  @parameterized.named_parameters(
+      ('negative_seconds', -5),
+      ('zero_seconds', 0),
+      ('negative_timedelta', datetime.timedelta(seconds=-1)),
+      ('zero_timedelta', datetime.timedelta(seconds=0)),
+  )
+  def test_time_scheduler_invalid_args(self, interval):
+    """TimeScheduler should raise an error if given invalid intervals."""
+    self.assertRaises(ValueError, schedulers.TimeScheduler, interval=interval)
+
+  def test_time_scheduler_logic(self):
+    """TimeScheduler should return True only after the interval has elapsed."""
+    current_time = 100.0
+
+    def mock_clock() -> float:
+      return current_time
+
+    scheduler = schedulers.TimeScheduler(interval=5, clock=mock_clock)
+
+    self.assertFalse(scheduler(None))
+
+    current_time = 104.0
+    self.assertFalse(scheduler(None))
+
+    current_time = 105.0
+    self.assertTrue(scheduler(None))
+
+    self.assertFalse(scheduler(None))
+
+    current_time = 111.0
+    self.assertTrue(scheduler(None))
+
+  def test_time_scheduler_timedelta_logic(self):
+    current_time = 100.0
+
+    def mock_clock() -> float:
+      return current_time
+
+    scheduler = schedulers.TimeScheduler(
+        interval=datetime.timedelta(seconds=5), clock=mock_clock
+    )
+
+    self.assertFalse(scheduler(None))
+
+    current_time = 105.0
+    self.assertTrue(scheduler(None))
+
+  def test_time_scheduler_boundary_conditions(self):
+    """TimeScheduler should trigger exactly at the boundary."""
+    current_time = 100.0
+
+    def mock_clock() -> float:
+      return current_time
+
+    scheduler = schedulers.TimeScheduler(interval=5.0, clock=mock_clock)
+    epsilon = 1e-6
+
+    # Verify T_0
+    self.assertFalse(scheduler(None))
+
+    # T_0 + interval - epsilon -> should not trigger
+    current_time = 100.0 + 5.0 - epsilon
+    self.assertFalse(scheduler(None))
+
+    # T_0 + interval -> should trigger
+    current_time = 100.0 + 5.0
+    self.assertTrue(scheduler(None))
+
+    # Immediately after -> should not trigger
+    self.assertFalse(scheduler(None))
+
+    # T_1 = 105.0
+    # T_1 + interval - epsilon -> should not trigger
+    current_time = 105.0 + 5.0 - epsilon
+    self.assertFalse(scheduler(None))
+
+    # T_1 + interval + epsilon -> should trigger
+    current_time = 105.0 + 5.0 + epsilon
+    self.assertTrue(scheduler(None))
+
+  def test_time_scheduler_timedelta_boundary_conditions(self):
+    """TimeScheduler should preserve microsecond precision with timedelta."""
+    current_time = 100.0
+
+    def mock_clock() -> float:
+      return current_time
+
+    interval = datetime.timedelta(microseconds=10)
+    scheduler = schedulers.TimeScheduler(interval=interval, clock=mock_clock)
+
+    # interval in seconds is 0.00001
+    interval_seconds = 0.00001
+    epsilon = 1e-6  # 1 microsecond
+
+    # Verify T_0
+    self.assertFalse(scheduler(None))
+
+    # T_0 + interval - epsilon -> should not trigger
+    current_time = 100.0 + interval_seconds - epsilon
+    self.assertFalse(scheduler(None))
+
+    # T_0 + interval -> should trigger
+    current_time = 100.0 + interval_seconds
+    self.assertTrue(scheduler(None))
+
+  def test_time_scheduler_large_time_jump(self):
+    """TimeScheduler should trigger once and reset baseline on large jump."""
+    current_time = 100.0
+
+    def mock_clock() -> float:
+      return current_time
+
+    scheduler = schedulers.TimeScheduler(interval=5.0, clock=mock_clock)
+
+    # Verify T_0
+    self.assertFalse(scheduler(None))
+
+    # Jump to 112.0 (elapsed 12.0)
+    current_time = 112.0
+    self.assertTrue(scheduler(None))  # Should trigger
+
+    # Subsequent call at same time should not trigger
+    self.assertFalse(scheduler(None))
+
+    # T_1 is now 112.0
+    # T_1 + 1.0 -> should not trigger
+    current_time = 113.0
+    self.assertFalse(scheduler(None))
+
+    # T_1 + 5.0 -> should trigger
+    current_time = 117.0
+    self.assertTrue(scheduler(None))
+
+  def test_time_scheduler_backward_clock(self):
+    """TimeScheduler should handle backward clock jumps gracefully."""
+    current_time = 100.0
+
+    def mock_clock() -> float:
+      return current_time
+
+    scheduler = schedulers.TimeScheduler(interval=5.0, clock=mock_clock)
+
+    # Verify T_0
+    self.assertFalse(scheduler(None))
+
+    # Move clock backward to 95.0
+    current_time = 95.0
+    self.assertFalse(scheduler(None))
+
+    # Move clock back to 100.0
+    current_time = 100.0
+    self.assertFalse(scheduler(None))
+
+    # Move clock to 105.0 (T_0 + 5.0) -> should trigger
+    current_time = 105.0
+    self.assertTrue(scheduler(None))
+
 
 if __name__ == '__main__':
   absltest.main()