fix: stop double-advancing the dataset cursor (half the dataset never trained)

tests/unit/test_trainer.py

@@ -80,6 +80,46 @@ def test_constructor_stores_config_and_initializes_attrs(cfg):
     assert trainer._server is None
 
 
+# ---------------------------------------------------------------------------
+# _pick_batch — dataset traversal
+# ---------------------------------------------------------------------------
+
+
+def test_pick_batch_covers_dataset_contiguously(tmp_path):
+    """Consecutive steps must train on consecutive rows, skipping none.
+
+    run() advances ``dataset_cursor`` by ``batch_size`` after every step, so
+    _pick_batch must rely on the cursor alone. Regression guard for the
+    double-advance bug where batches started at 0, 2b, 4b, … and every other
+    slice was never trained on.
+    """
+    cfg = _make_cfg(tmp_path, training={"batch_size": 2})
+    trainer = TinkerNeMoGymTrainer(cfg)
+    trainer.dataset = [{"i": k} for k in range(100)]
+
+    covered = []
+    for _ in range(5):
+        covered.extend(row["i"] for row in trainer._pick_batch())
+        trainer.dataset_cursor += cfg.training.batch_size  # mirrors run()
+
+    assert covered == list(range(10))
+
+
+def test_pick_batch_stays_contiguous_after_resume(tmp_path):
+    """A restored cursor continues contiguously from where it left off."""
+    cfg = _make_cfg(tmp_path, training={"batch_size": 2})
+    trainer = TinkerNeMoGymTrainer(cfg)
+    trainer.dataset = [{"i": k} for k in range(100)]
+    trainer.dataset_cursor = 20  # restored from checkpoint meta
+
+    covered = []
+    for _ in range(3):
+        covered.extend(row["i"] for row in trainer._pick_batch())
+        trainer.dataset_cursor += cfg.training.batch_size
+
+    assert covered == [20, 21, 22, 23, 24, 25]
+
+
 # ---------------------------------------------------------------------------
 # setup()
 # ---------------------------------------------------------------------------

tinker_nemogym/trainer.py

@@ -517,7 +517,7 @@ async def _run_step(self, step: int) -> dict:
         if not self.dataset:
             raise RuntimeError("dataset is empty; did setup() run?")
 
-        batch_prompts = self._pick_batch(step)
+        batch_prompts = self._pick_batch()
         groups_raw = await self._do_rollouts(step, batch_prompts)
         if groups_raw is None:
             return {"step": step, "mean_reward": 0.0, "n_datums": 0, "n_dropped": 0}
@@ -539,16 +539,21 @@ async def _run_step(self, step: int) -> dict:
     # _run_step helpers
     # ------------------------------------------------------------------
 
-    def _pick_batch(self, step: int) -> list[dict]:
+    def _pick_batch(self) -> list[dict]:
         """Slice ``batch_size`` prompts from the dataset with wrap-around.
 
-        When :attr:`cfg.training.random_seed` is set and we're resuming,
-        ``dataset_cursor`` is used as the starting offset so we don't
-        re-train on the same prefix.
+        ``dataset_cursor`` is the sole start offset: ``run()`` advances it by
+        ``batch_size`` after every step, and it is restored from the checkpoint
+        on resume, so it already points at the next unseen batch.
+
+        It previously *also* added ``step * batch_size`` on top of the cursor.
+        Because the cursor is advanced independently each step, that double-
+        counted the advance — batches started at 0, 2*batch_size, 4*batch_size,
+        … and every other slice of the dataset was silently never trained on
+        (and the offset was wrong again after a resume).
         """
         batch_size = self.cfg.training.batch_size
-        cursor_base = self.dataset_cursor
-        start = (cursor_base + step * batch_size) % len(self.dataset)
+        start = self.dataset_cursor % len(self.dataset)
         batch = [self.dataset[(start + i) % len(self.dataset)] for i in range(batch_size)]
         return batch