From f321a1622728df7d25720c950a0d928cd4aef141 Mon Sep 17 00:00:00 2001 From: Tony Li Date: Fri, 10 Jul 2026 11:01:25 -0700 Subject: [PATCH] Replace default BernoulliStepScheduler with TimeScheduler (defaulting to 10 minutes) for deterministic flushing. PiperOrigin-RevId: 945797039 --- envlogger/backends/riegeli_backend_writer.py | 9 +- .../backends/riegeli_backend_writer_test.py | 24 +++ envlogger/backends/schedulers.py | 32 ++++ envlogger/backends/schedulers_test.py | 159 +++++++++++++++++- 4 files changed, 220 insertions(+), 4 deletions(-) diff --git a/envlogger/backends/riegeli_backend_writer.py b/envlogger/backends/riegeli_backend_writer.py index 42824a9..461a791 100644 --- a/envlogger/backends/riegeli_backend_writer.py +++ b/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() diff --git a/envlogger/backends/riegeli_backend_writer_test.py b/envlogger/backends/riegeli_backend_writer_test.py index 716d5a8..1e34a9a 100644 --- a/envlogger/backends/riegeli_backend_writer_test.py +++ b/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() diff --git a/envlogger/backends/schedulers.py b/envlogger/backends/schedulers.py index 05241d7..b067035 100644 --- a/envlogger/backends/schedulers.py +++ b/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 diff --git a/envlogger/backends/schedulers_test.py b/envlogger/backends/schedulers_test.py index 538728c..bdf61b6 100644 --- a/envlogger/backends/schedulers_test.py +++ b/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()