`RLLearner` resume resets the data-shuffle PRNG state

[yixiaoer]

I first found this while testing GRPO checkpoint/resume on TPU. After reducing the failure, the underlying issue is in the base RLLearner and can be reproduced on CPU without running a model or training.

Expected Behavior
After a fresh-process resume, the learner should continue the same data-shuffle PRNG stream as an uninterrupted run:

shuffle_state(resumed) == shuffle_state(uninterrupted)   # at the same resume point

Actual Behavior

RLLearner initializes _data_shuffle_seed from data_shuffle_seed and advances it on every shuffle:

• initialization:

  tunix/tunix/rl/rl_learner.py

  Lines 101 to 105 in caa444b

  	self._data_shuffle_seed = (
  	jax.random.PRNGKey(data_shuffle_seed)
  	if data_shuffle_seed is not None
  	else None
  	)

• advancement:
  

  tunix/tunix/rl/rl_learner.py

  Lines 347 to 354 in caa444b

  	for _ in range(repeats):
  	if self._data_shuffle_seed is not None:
  	shuffle_seed, self._data_shuffle_seed = jax.random.split(
  	self._data_shuffle_seed
  	)
  	shuffled_indices = jax.random.permutation(shuffle_seed, len(examples))
  	for i in shuffled_indices:
  	data_queue.put([examples[i]])

On resume, trainer progress is restored and _prepare_data fast-forwards the input iterator, but the _data_shuffle_seed is not restored:

• checkpoint metadata:
  

  tunix/tunix/rl/rl_cluster.py

  Lines 469 to 472 in caa444b

  	custom_checkpoint_metadata_fn=lambda: {
  	"global_step": self.global_steps + 1,
  	"role": Role.ACTOR.value,
  	}, # offset by 1 since global_step is incremented after the training loop in rl_learner. # pylint: disable=line-too-long

• iterator fast-forward:
  

  tunix/tunix/rl/rl_learner.py

  Lines 398 to 407 in caa444b

  	try:
  	while True:
  	while (
  	mode == rl_engine_lib.Mode.TRAIN
  	and self._iter_steps < self._last_iter_step
  	): # fast forward the iterator if loading from a previous checkpoint.
  	next(iterator)
  	self._iter_steps += 1
  	if self._iter_steps == self._last_iter_step:
  	logging.info("Fast forwarded %d micro-batches.", self._iter_steps)

So the resumed learner reaches the same post-checkpoint data in this reproduction, but its shuffle state starts again from PRNGKey(data_shuffle_seed).

Steps to Reproduce the Problem

This test uses the real RLLearner.__init__ and _prepare_data. num_iterations=1, so repeated GRPO updates are not involved.

Save as test_resume_shuffle_state.py:

from unittest import mock

from flax import nnx
import numpy as np
import optax

from tunix.rl import algorithm_config, rl_cluster, rl_learner
from tunix.rl.queue import data_queue


N = 8
DATA = [f"{g}{i}" for g in "AB" for i in range(N)]


class Learner(rl_learner.RLLearner):
  _generate_and_compute_advantage = lambda self, x, mode: x
  _compute_trajectory_ids = lambda self, x, steps: [""] * len(x["prompts"])
  _num_iterations = lambda self: 1
  _num_generations = lambda self: 1


def make_learner(iter_steps=0, global_step=0):
  engine = mock.MagicMock()
  engine.cluster_config.training_config = rl_cluster.RLTrainingConfig(
      actor_optimizer=optax.sgd(0.1),
      eval_every_n_steps=1,
      max_steps=2,
  )
  engine.actor_trainer.model = nnx.Module()
  engine.rollout.model.return_value = nnx.Module()
  engine.actor_trainer.iter_steps = iter_steps
  engine.actor_trainer.restored_global_step.return_value = global_step

  return Learner(
      engine,
      algorithm_config.AlgorithmConfig(),
      reward_fns=lambda **kw: [],
      data_shuffle_seed=0,
  )


def prepare(learner, iterator):
  q = data_queue.SimpleDataQueue(maxsize=0)
  learner._prepare_data(
      iterator,
      proceed_num_steps=N,
      sample_repeat=1,
      batch_repeat=1,
      service_target_batch_size=1,
      data_queue=q,
  )
  return [str(x[0]["prompts"][0]) for x in iter(q.get, None)]


def test_resume_preserves_shuffle_state():
  batches = lambda: iter(
      [{"prompts": np.array([x])} for x in DATA]
  )

  continuous = make_learner()
  iterator = batches()

  prepare(continuous, iterator)  # group A
  expected_state = np.asarray(
      continuous._data_shuffle_seed
  ).copy()
  expected = prepare(continuous, iterator)  # group B

  resumed = make_learner(iter_steps=N, global_step=1)
  actual_state = np.asarray(
      resumed._data_shuffle_seed
  ).copy()
  actual = prepare(resumed, batches())  # skips A, processes B

  print("expected shuffle state:", expected_state)
  print("resumed shuffle state: ", actual_state)
  print("same post-checkpoint data:", sorted(actual) == sorted(expected))

  assert sorted(actual) == sorted(expected) == DATA[N:]
  np.testing.assert_array_equal(actual_state, expected_state)

Run:

JAX_PLATFORMS=cpu python -m pytest -q -s test_resume_shuffle_state.py

Observed output:

expected shuffle state: [ 928981903 3453687069]
resumed shuffle state:  [0 0]
same post-checkpoint data: True

E     AssertionError:
E     Arrays are not equal
E      ACTUAL: array([0, 0], dtype=uint32)
E      DESIRED: array([ 928981903, 3453687069], dtype=uint32)

Both learners reach the same post-checkpoint data (group B), but the resumed learner has restarted the shuffle PRNG state.

I also encountered the same behavior through the real Trainer + Orbax save/restore path in separate processes.

Trainer progress is restored and the input advances to the same post-checkpoint data, while _data_shuffle_seed is reset to the initial seed. Injecting only the expected post-checkpoint shuffle state makes the resumed shuffle key and permutation match the uninterrupted run.

Environment

• CPU reproducer (JAX_PLATFORMS=cpu)
• originally observed on TPU during GRPO checkpoint/resume testing
• google-tunix 0.1.8, Tunix main @ caa444b73ffb1eb28b706f82c152fbe6ad7d5cd3
• JAX 0.11.0
• Flax 0.12.8
• Python 3.12

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided all the required information in the "Environment" section.
• I have provided a minimal, reproducible example.

Would you like to help us fix it?
Yes.