[Orbit] Add PeriodicAction utility

orbit/actions/__init__.py

@@ -74,3 +74,5 @@ class to make it easy to trigger actions conditionally based on reusable
 from orbit.actions.new_best_metric import NewBestMetric
 
 from orbit.actions.save_checkpoint_if_preempted import SaveCheckpointIfPreempted
+
+from orbit.actions.periodic_action import PeriodicAction

orbit/actions/periodic_action.py

@@ -0,0 +1,44 @@
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Provides the `PeriodicAction` abstraction."""
+
+from typing import Callable
+
+from orbit import runner
+import tensorflow as tf
+
+
+class PeriodicAction:
+  """Wraps an action to be executed only at a specific step interval."""
+
+  def __init__(self,
+               action: Callable[[runner.Output], None],
+               interval: int,
+               global_step: tf.Variable):
+    """Initializes the instance.
+
+    Args:
+      action: The action (callable) to wrap.
+      interval: The interval (in global steps) at which to execute the action.
+      global_step: The global step variable.
+    """
+    self._action = action
+    self._interval = interval
+    self._global_step = global_step
+
+  def __call__(self, output: runner.Output) -> None:
+    # Execute action only if the current step is divisible by the interval.
+    if self._global_step.numpy() % self._interval == 0:
+      self._action(output)

orbit/actions/periodic_action_test.py

@@ -0,0 +1,58 @@
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for orbit.actions.periodic_action."""
+
+from orbit import actions
+from orbit.utils import common
+import tensorflow as tf
+
+
+class PeriodicActionTest(tf.test.TestCase):
+
+  def test_periodic_execution(self):
+    global_step = common.create_global_step()
+    call_count = 0
+
+    def action(_):
+      nonlocal call_count
+      call_count += 1
+
+    # Execute action every 10 steps.
+    periodic_action = actions.PeriodicAction(
+        action=action, interval=10, global_step=global_step)
+
+    # Step 5: Should not execute.
+    global_step.assign(5)
+    periodic_action({})
+    self.assertEqual(call_count, 0)
+
+    # Step 10: Should execute.
+    global_step.assign(10)
+    periodic_action({})
+    self.assertEqual(call_count, 1)
+
+    # Step 15: Should not execute.
+    global_step.assign(15)
+    periodic_action({})
+    self.assertEqual(call_count, 1)
+
+    # Step 20: Should execute again.
+    global_step.assign(20)
+    periodic_action({})
+    self.assertEqual(call_count, 2)
+
+
+if __name__ == '__main__':
+  tf.test.main()