diff --git a/tests/rl/agentic/agentic_grpo_learner_test.py b/tests/rl/agentic/agentic_grpo_learner_test.py index 0304039bd..9bc302b3b 100644 --- a/tests/rl/agentic/agentic_grpo_learner_test.py +++ b/tests/rl/agentic/agentic_grpo_learner_test.py @@ -109,10 +109,13 @@ def _mock_generate( else: prompt_tokens.append(tokenizer.encode(" ".join(m["content"] for m in p))) max_p_len = max(len(pt) for pt in prompt_tokens) - padded_prompts = np.array([ - np.pad(pt, (max(0, max_p_len - len(pt)), 0), constant_values=0) - for pt in prompt_tokens - ], dtype=np.int32) + padded_prompts = np.array( + [ + np.pad(pt, (max(0, max_p_len - len(pt)), 0), constant_values=0) + for pt in prompt_tokens + ], + dtype=np.int32, + ) return base_rollout.RolloutOutput( text=text, tokens=tokens, @@ -585,9 +588,7 @@ def test_compute_logps_chunk_size(self): chat_parser=MockChatParser(), ) - train_ds = _dummy_dataset( - MySource(data=["1", "2"], repeat=1), batch_size=2 - ) + train_ds = _dummy_dataset(MySource(data=["1", "2"], repeat=1), batch_size=2) with ( mock.patch.object( @@ -639,9 +640,7 @@ class MockModel(nnx.Module): def __init__(self, *, rngs: nnx.Rngs): self.lm_head = 1 - def __call__( - self, inputs, positions, cache, attention_mask, **kwargs - ): + def __call__(self, inputs, positions, cache, attention_mask, **kwargs): del kwargs return ( jnp.full( @@ -703,9 +702,7 @@ class MockModel(nnx.Module): def __init__(self, *, rngs: nnx.Rngs): self.lm_head = 1 - def __call__( - self, inputs, positions, cache, attention_mask, **kwargs - ): + def __call__(self, inputs, positions, cache, attention_mask, **kwargs): del kwargs return ( jnp.full( @@ -2375,6 +2372,122 @@ def _patch_process_results( decoded_completion.count("Assistant:"), 2 ) # 3 turns but terminal env obs does not append generation msg + def test_force_on_policy_ratio_bypasses_actor_recompute(self): + """Verifies force_on_policy_ratio=True sets old_logps to None with 0 extra passes.""" + vocab = _mock_vocab() + tokenizer = tokenizer_adapter.TokenizerAdapter(vocab) + model = test_common.ToyTransformer( + config=test_common.ModelConfig(vocab_size=vocab.GetPieceSize()), + rngs=nnx.Rngs(0), + ) + ref_model = test_common.ToyTransformer( + config=test_common.ModelConfig(vocab_size=vocab.GetPieceSize()), + rngs=nnx.Rngs(0), + ) + mesh = pxla.thread_resources.env.physical_mesh + cluster_config = rl_engine_lib.ClusterConfig( + role_to_mesh={ + rl_engine_lib.Role.ACTOR: mesh, + rl_engine_lib.Role.REFERENCE: mesh, + rl_engine_lib.Role.ROLLOUT: mesh, + }, + rollout_engine="vanilla", + offload_to_cpu=False, + training_config=rl_engine_lib.RLTrainingConfig( + actor_optimizer=optax.sgd(1e-3), + eval_every_n_steps=10, + max_steps=10, + ), + rollout_config=base_rollout.RolloutConfig( + max_prompt_length=32, + max_tokens_to_generate=10, + return_logprobs=True, + ), + ) + rl_engine = rl_engine_lib.RLEngine( + actor=model, + reference=ref_model, + tokenizer=tokenizer, + cluster_config=cluster_config, + ) + grpo_config = agentic_grpo_learner.GRPOConfig( + beta=0.0, + force_compute_kl=False, + max_response_length=10, + num_generations=2, + num_iterations=1, + use_rollout_logps=True, + force_on_policy_ratio=True, + ) + learner = agentic_grpo_learner.GRPOLearner( + rl_engine=rl_engine, + reward_fns=reward_fn_1, + algo_config=grpo_config, + chat_parser=MockChatParser(), + ) + + class MockTraj: + + def __init__(self, index): + self.traj = { + "conversation_text": [ + {"role": "assistant", "content": f"msg {index}"} + ], + "conversation_tokens": np.array([1, 2, 3]), + "conversation_masks": np.array([1, 1, 1]), + "old_logprobs": np.full(3, 1.0, dtype=np.float32), + "policy_version": 0, + "trajectory_reward": 1.0, + "prompt_tokens": np.array([4, 5]), + "original_input": {"prompts": "hello"}, + "group_id": "test_group", + } + + trajectories = [MockTraj(0), MockTraj(1)] + + with mock.patch.object( + rl_engine, + "get_actor_per_token_logps", + return_value=jnp.full((2, 10), -1.0), + autospec=True, + ) as mock_get_actor_logps: + results = learner._process_results(trajectories, expected_step=1) + self.assertLen(results, 1) + train_example = results[0] + + # 1. Asserts 0 extra trainer forward passes! + mock_get_actor_logps.assert_not_called() + + # 2. Asserts old_per_token_logps is None (forces ratio to 1.0 via stop_gradient in loss) + self.assertIsNone(train_example.old_per_token_logps) + + def test_force_on_policy_ratio_config_validation(self): + """force_on_policy_ratio rejects multi-iteration, allows stale trajectories.""" + # num_iterations > 1 raises ValueError: old_logp is re-derived from the current + # policy on every inner epoch, so the trust region vanishes after the first. + with self.assertRaisesRegex( + ValueError, "can only be True when num_iterations == 1" + ): + agentic_grpo_learner.GRPOConfig( + num_generations=2, + num_iterations=2, + force_on_policy_ratio=True, + ) + + # off_policy_steps > 0 is a supported trade-off (near-on-policy training on + # slightly stale trajectories), so it warns rather than raising. + with mock.patch.object( + agentic_grpo_learner.logging, "warning" + ) as mock_warn: + config = agentic_grpo_learner.GRPOConfig( + num_generations=2, + num_iterations=1, + off_policy_steps=1, + force_on_policy_ratio=True, + ) + mock_warn.assert_called_once() + self.assertIn("off_policy_steps", mock_warn.call_args[0][0]) + if __name__ == "__main__": absltest.main() diff --git a/tunix/rl/agentic/agentic_grpo_learner.py b/tunix/rl/agentic/agentic_grpo_learner.py index 15ca33640..c477023b1 100644 --- a/tunix/rl/agentic/agentic_grpo_learner.py +++ b/tunix/rl/agentic/agentic_grpo_learner.py @@ -37,19 +37,19 @@ import jax import jax.numpy as jnp import numpy as np -from tunix.rl import algo_core # pylint: disable=unused-import from tunix.perf.experimental import constants as perf_constants +from tunix.rl import algo_core # pylint: disable=unused-import from tunix.rl import common from tunix.rl import function_registry from tunix.rl import rl_cluster as rl_engine_lib from tunix.rl import utils as rl_utils from tunix.rl.agentic import agentic_rl_learner from tunix.rl.agentic import utils as agentic_utils -from tunix.utils import compat from tunix.rl.agentic.agents import base_agent from tunix.rl.agentic.agents import model_agent from tunix.rl.agentic.environments import base_environment from tunix.rl.agentic.environments import task_environment +from tunix.utils import compat from tunix.utils import trajectory_logger TrainingInputT = agentic_rl_learner.TrainingInputT @@ -83,6 +83,17 @@ class GRPOConfig(agentic_rl_learner.AgenticRLConfig): max_concurrency: Maximum number of concurrent rollout engines. off_policy_steps: Number of off-policy steps can be accepted before a policy update. + use_rollout_logps: Use the rollout engine's log-probabilities as + old_per_token_logps. False makes the trainer recompute them. + force_on_policy_ratio: When num_iterations == 1, use + stop_gradients(current_logp) as old_per_token_logps instead of recomputing + or using rollout logps. Pins the surrogate ratio to 1.0, so clipping never + fires and sampler-vs-trainer numerical noise is removed from the ratio. + log_sampler_trainer_agreement: Optionally spend one extra trainer forward + pass per step to log sampler-vs-trainer log-probability agreement metrics. + Without force_on_policy_ratio these metrics come for free from the logps + already being computed; with it, no trainer logps exist, so this flag + pays for them explicitly. Default False degenerate_group_masking: Whether to mask out degenerate groups with all-0 advantages. Deprecated. Will remove in the next release. """ @@ -112,6 +123,12 @@ class GRPOConfig(agentic_rl_learner.AgenticRLConfig): False # Whether to mask out degenerate groups with all-0 advantages. ) use_rollout_logps: bool = True + # Pin the surrogate ratio to 1.0 (old_logp := stop_gradient(current_logp)). + # Valid for single-iteration on-policy training only. + force_on_policy_ratio: bool = False + # Costs one trainer forward pass; keeps the sampler/trainer agreement metrics + # alive when force_on_policy_ratio would otherwise leave nothing to compare. + log_sampler_trainer_agreement: bool = False # Truncated importance-sampling (TIS) correction for the residual mismatch # between the rollout sampler and the trainer's recomputed log-probabilities. # Set to ``"token"`` to enable per-token TIS weights. When enabled, the loss @@ -140,6 +157,28 @@ def __post_init__(self): "loss_algo should be either grpo or gspo-token. Received: " f"{self.loss_algo}" ) + if self.force_on_policy_ratio: + if self.num_iterations > 1: + raise ValueError( + "force_on_policy_ratio can only be True when num_iterations == 1." + " With num_iterations > 1 the policy is updated several times on" + " the same batch, so the surrogate ratio is genuinely != 1 after" + " the first inner epoch; pinning it to 1 removes the trust region" + " for every subsequent epoch. Got" + f" num_iterations={self.num_iterations}" + ) + + if self.off_policy_steps > 0: + logging.warning( + "force_on_policy_ratio=True with off_policy_steps=%d: trajectories " + "may be up to %d policy updates stale, but the surrogate ratio is " + "pinned to 1.0, so the off-policy correction is discarded. This is " + "a deliberate trade for near-on-policy training; pair it with an " + "importance-sampling correction if the behavior and target " + "policies can drift apart.", + self.off_policy_steps, + self.off_policy_steps, + ) TGrpoConfig = TypeVar("TGrpoConfig", bound=GRPOConfig) @@ -228,9 +267,11 @@ def __init__( else: logging.warning("Metrics log dir is None, skipping trajectory logging.") - self.algo_config.temperature = self.rl_engine.get_rollout_config( # pyrefly: ignore[missing-attribute] - mode=rl_engine_lib.Mode.TRAIN - ).temperature + self.algo_config.temperature = ( + self.rl_engine.get_rollout_config( # pyrefly: ignore[missing-attribute] + mode=rl_engine_lib.Mode.TRAIN + ).temperature + ) # Workaround to pass loss fn with algorithm flag policy_loss_fn = function_registry.get_policy_loss_fn( @@ -414,6 +455,7 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample: if ( example.old_per_token_logps is None and not self.algo_config.use_rollout_logps + and not self.algo_config.force_on_policy_ratio ): updates["old_per_token_logps"] = self.rl_engine.get_actor_per_token_logps( prompt_tokens=prompt_tokens, @@ -439,7 +481,11 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample: # The rollout-logps path defers its trainer recompute here too. Not just # diagnostics: sampler_is="token" consumes it as old_per_token_logps. need_trainer_logps = ( - self.algo_config.use_rollout_logps + ( + not self.algo_config.force_on_policy_ratio + or self.algo_config.log_sampler_trainer_agreement + ) + and self.algo_config.use_rollout_logps and example.old_per_token_logps is not None and (self._have_actor_mesh() or self.algo_config.sampler_is == "token") ) @@ -465,7 +511,10 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample: ) if sampler_is_weights is not None: updates["sampler_is_weights"] = sampler_is_weights - if self.algo_config.sampler_is == "token": + if ( + self.algo_config.sampler_is == "token" + and not self.algo_config.force_on_policy_ratio + ): updates["old_per_token_logps"] = trainer_logps if updates: @@ -573,7 +622,7 @@ def _process_results( ) if ( len(completion_tokens) >= max_response_length - and completion_mask[-1] != eos_value + and completion_tokens[-1] != eos_value ): clipped_completion_count += 1 padded_prompt, padded_completion, _ = ( @@ -646,7 +695,24 @@ def _process_results( rollout_per_token_logps = None trainer_per_token_logps = None - if self.algo_config.use_rollout_logps and padded_old_logprobs: + if self.algo_config.force_on_policy_ratio: + # The loss derives old_per_token_logps from the actor's own forward pass. + old_per_token_logps = None + if padded_old_logprobs: + rollout_per_token_logps = jnp.asarray(padded_old_logprobs) + if ( + self.algo_config.log_sampler_trainer_agreement + and have_actor_mesh + and not is_packed + ): + trainer_per_token_logps = self.rl_engine.get_actor_per_token_logps( + prompt_tokens=prompt_ids, + completion_tokens=completion_ids, + pad_id=pad_value, + eos_id=eos_value, + micro_batch_size=compute_logps_micro_batch_size, + ) + elif self.algo_config.use_rollout_logps and padded_old_logprobs: rollout_per_token_logps = jnp.asarray(padded_old_logprobs) old_per_token_logps = rollout_per_token_logps # The diagnostic pass (and the sampler-IS ``token`` path, which needs the @@ -845,11 +911,11 @@ def _process_results( f"{prefix}/{sub_key}/max": (np.max(flat_vals), np.max), f"{prefix}/{sub_key}/min": (np.min(flat_vals), np.min), }) - self.rl_engine.buffer_metrics_async( - metrics_to_log, # pyrefly: ignore[bad-argument-type] - mode=mode, - step=expected_step, # pyrefly: ignore[bad-argument-type] - ) + self.rl_engine.buffer_metrics_async( + metrics_to_log, # pyrefly: ignore[bad-argument-type] + mode=mode, + step=expected_step, # pyrefly: ignore[bad-argument-type] + ) for metric_fn in self.metric_fns: user_defined_metric = metric_fn( @@ -864,7 +930,9 @@ def _process_results( }, ) self.rl_engine.buffer_metrics_async( - user_defined_metric, mode=mode, step=expected_step # pyrefly: ignore[bad-argument-type] + user_defined_metric, + mode=mode, + step=expected_step, # pyrefly: ignore[bad-argument-type] ) combined_batch = TrainExample(