diff --git a/reagent/gym/datasets/replay_buffer_dataset.py b/reagent/gym/datasets/replay_buffer_dataset.py
index 053aeb2ba..eea3c713b 100644
--- a/reagent/gym/datasets/replay_buffer_dataset.py
+++ b/reagent/gym/datasets/replay_buffer_dataset.py
@@ -134,6 +134,8 @@ def __iter__(self):
                 possible_actions_mask = next_possible_actions_mask
                 num_steps += 1
                 global_num_steps += 1
+                if self._agent.post_step:
+                    self._agent.post_step(transition)
             if self._post_episode_callback:
                 self._post_episode_callback(trajectory, info)