From 8dbb0b02994ec384d386550eb40b5760645ee94f Mon Sep 17 00:00:00 2001 From: Kingpython123 Date: Sat, 31 Oct 2026 09:18:21 -0700 Subject: [PATCH] Add repetition_filter operator for degenerate rollout responses (#620) Add a RepetitionFilter experience operator that drops experiences whose generated response has collapsed into repetition. The response is scored by its duplicate n-gram ratio (1 - unique_ngrams / total_ngrams) over the model-generated tokens only, ignoring prompt tokens and tokens masked out by action_mask. Experiences are dropped when the ratio is strictly greater than max_repeat_ratio; empty and very short responses are always kept. - Register as repetition_filter in EXPERIENCE_OPERATORS - Add unit tests for the ratio function, the filter, action_mask handling and registry creation - Document built-in operators in the operator development guide (en/zh) --- .../source/tutorial/develop_operator.md | 16 ++ .../source_zh/tutorial/develop_operator.md | 16 ++ tests/buffer/repetition_filter_test.py | 164 ++++++++++++++++++ trinity/buffer/operators/__init__.py | 1 + .../operators/filters/repetition_filter.py | 121 +++++++++++++ 5 files changed, 318 insertions(+) create mode 100644 tests/buffer/repetition_filter_test.py create mode 100644 trinity/buffer/operators/filters/repetition_filter.py diff --git a/docs/sphinx_doc/source/tutorial/develop_operator.md b/docs/sphinx_doc/source/tutorial/develop_operator.md index 5f2d62b0575..49a8b99a7ff 100644 --- a/docs/sphinx_doc/source/tutorial/develop_operator.md +++ b/docs/sphinx_doc/source/tutorial/develop_operator.md @@ -81,6 +81,22 @@ The `RewardFilter` reduces the number of experiences, which may cause the traine The above setting means that the `Explorer` will sync with the `Trainer` every 2 steps and will continue running regardless of how many steps the `Trainer` has completed. This ensures that the `Trainer` can always get enough experiences to start a training step as long as the `Explorer` is running. ``` +### Built-in Operators + +Trinity-RFT ships with the following experience operators, all registered in {class}`trinity.buffer.operators.EXPERIENCE_OPERATORS` and usable by name in `experience_pipeline.operators`: + +| Name | Description | Main Arguments | +|------|-------------|----------------| +| `reward_filter` | Drop experiences whose reward is below a threshold. | `threshold` (default `0.0`) | +| `reward_std_filter` | Drop task groups whose reward standard deviation is not above a threshold (no learning signal). | `threshold` (default `0.0`) | +| `dapo_dynamic_sampling` | DAPO dynamic sampling: keep a task group only when some but not all rollouts are correct. | `metric_key` (default `accuracy`), `correct_threshold` (default `0.0`) | +| `mask_response_truncated` | Zero the `action_mask` of truncated responses so they do not contribute to the loss. | - | +| `invalid_reward_filter` | Drop experiences whose reward is `None` or `NaN`. | - | +| `repetition_filter` | Drop experiences whose generated response is highly repetitive (repetition collapse). The response is scored by its duplicate n-gram ratio `1 - unique_ngrams / total_ngrams` over the model-generated tokens only (prompt tokens and tokens with a zero `action_mask` are ignored), and dropped when the ratio is strictly greater than `max_repeat_ratio`. Empty responses and responses shorter than `min_response_tokens` are always kept. | `ngram_size` (default `4`), `max_repeat_ratio` (default `0.5`), `min_response_tokens` (default `32`) | +| `reward_shaping_mapper` | Reshape rewards using statistics stored in `Experience.info`. | `reward_shaping_configs` | +| `pass_rate_calculator` | Compute the pass rate of each task group. | - | +| `data_juicer` | Apply Data-Juicer operators to experiences. | `config_path` | + ### Advanced Features #### Using Auxiliary Models inside Operators diff --git a/docs/sphinx_doc/source_zh/tutorial/develop_operator.md b/docs/sphinx_doc/source_zh/tutorial/develop_operator.md index 5c9e1acded7..4084f88fefc 100644 --- a/docs/sphinx_doc/source_zh/tutorial/develop_operator.md +++ b/docs/sphinx_doc/source_zh/tutorial/develop_operator.md @@ -90,6 +90,22 @@ synchronizer: 上述设置意味着 `Explorer` 每运行 2 步就会与 `Trainer` 同步一次，且无论 `Trainer` 当前完成了多少步都会继续运行。这确保了只要 `Explorer` 在运行，`Trainer` 就总能获得足够的 experience 来启动训练步骤。 ``` +### 内置算子 + +Trinity-RFT 内置了以下 experience 算子，它们均已在 {class}`trinity.buffer.operators.EXPERIENCE_OPERATORS` 中注册，可在 `experience_pipeline.operators` 中通过名称直接使用： + +| 名称 | 说明 | 主要参数 | +|------|------|----------| +| `reward_filter` | 丢弃 reward 低于阈值的 experience。 | `threshold`（默认 `0.0`） | +| `reward_std_filter` | 丢弃组内 reward 标准差不高于阈值（没有学习信号）的任务组。 | `threshold`（默认 `0.0`） | +| `dapo_dynamic_sampling` | DAPO 动态采样：仅保留部分（而非全部）rollout 正确的任务组。 | `metric_key`（默认 `accuracy`）、`correct_threshold`（默认 `0.0`） | +| `mask_response_truncated` | 将被截断的响应的 `action_mask` 置零，使其不参与 loss 计算。 | - | +| `invalid_reward_filter` | 丢弃 reward 为 `None` 或 `NaN` 的 experience。 | - | +| `repetition_filter` | 丢弃生成响应高度重复（重复坍缩）的 experience。仅对模型生成的 token 计算重复 n-gram 比例 `1 - unique_ngrams / total_ngrams`（忽略 prompt token 以及 `action_mask` 为 0 的 token），当该比例严格大于 `max_repeat_ratio` 时丢弃。空响应及短于 `min_response_tokens` 的响应总是保留。 | `ngram_size`（默认 `4`）、`max_repeat_ratio`（默认 `0.5`）、`min_response_tokens`（默认 `32`） | +| `reward_shaping_mapper` | 利用 `Experience.info` 中的统计量对 reward 进行重塑。 | `reward_shaping_configs` | +| `pass_rate_calculator` | 计算每个任务组的通过率。 | - | +| `data_juicer` | 对 experience 应用 Data-Juicer 算子。 | `config_path` | + ### 进阶特性 #### 在 Operator 中使用辅助模型 diff --git a/tests/buffer/repetition_filter_test.py b/tests/buffer/repetition_filter_test.py new file mode 100644 index 00000000000..3ca384c68cc --- /dev/null +++ b/tests/buffer/repetition_filter_test.py @@ -0,0 +1,164 @@ +import unittest +from typing import List, Optional + +import torch + +from trinity.buffer.operators.experience_operator import create_operators +from trinity.buffer.operators.filters.repetition_filter import ( + RepetitionFilter, + compute_repeat_ratio, +) +from trinity.common.config import OperatorConfig +from trinity.common.experience import EID, Experience + +PROMPT = [1, 2, 3, 4] + + +def make_exp( + response: List[int], + task: int = 0, + run: int = 0, + action_mask: Optional[List[int]] = None, +) -> Experience: + """Build an experience with a fixed prompt and the given response tokens.""" + return Experience( + eid=EID(task=task, run=run), + tokens=torch.tensor(PROMPT + response, dtype=torch.long), + prompt_length=len(PROMPT), + reward=1.0, + action_mask=None if action_mask is None else torch.tensor(action_mask, dtype=torch.bool), + ) + + +def diverse(length: int) -> List[int]: + return list(range(100, 100 + length)) + + +def repetitive(length: int) -> List[int]: + pattern = [7, 8, 9] + return (pattern * (length // len(pattern) + 1))[:length] + + +class TestComputeRepeatRatio(unittest.TestCase): + def test_diverse_sequence_has_zero_ratio(self): + self.assertEqual(compute_repeat_ratio(diverse(40), 4), 0.0) + + def test_repetitive_sequence_has_high_ratio(self): + # 60 tokens -> 57 4-grams, only 3 of them unique + ratio = compute_repeat_ratio([7, 8, 9] * 20, 4) + self.assertAlmostEqual(ratio, 1 - 3 / 57) + + def test_sequence_shorter_than_ngram_is_zero(self): + self.assertEqual(compute_repeat_ratio([1, 2], 4), 0.0) + self.assertEqual(compute_repeat_ratio([], 4), 0.0) + + def test_single_duplicate(self): + # 5 tokens -> 4 bigrams: (1,1) (1,1) (1,2) (2,3) -> 3 unique + self.assertAlmostEqual(compute_repeat_ratio([1, 1, 1, 2, 3], 2), 0.25) + + +class TestRepetitionFilter(unittest.IsolatedAsyncioTestCase): + async def test_keeps_diverse_and_drops_repetitive(self): + op = RepetitionFilter(ngram_size=4, max_repeat_ratio=0.5, min_response_tokens=32) + exps = [ + make_exp(diverse(64), task=0), + make_exp(repetitive(64), task=1), + make_exp(diverse(64), task=2), + ] + kept, metrics = await op.process(exps) + self.assertEqual([e.eid.task for e in kept], [0, 2]) + self.assertEqual(metrics["filtered_count"], 1) + self.assertEqual(metrics["skipped_empty_count"], 0) + self.assertEqual(metrics["skipped_short_count"], 0) + self.assertGreater(metrics["repeat_ratio/max"], 0.9) + self.assertAlmostEqual(metrics["repeat_ratio/mean"], metrics["repeat_ratio/max"] / 3) + + async def test_empty_and_short_responses_are_kept(self): + op = RepetitionFilter(ngram_size=4, max_repeat_ratio=0.5, min_response_tokens=32) + exps = [ + make_exp([], action_mask=[]), # multi-turn experience with no response tokens + make_exp(repetitive(64), action_mask=[0] * 64), # all tokens are observations + make_exp(repetitive(8)), # repetitive but too short to judge + make_exp([5]), # single token + ] + kept, metrics = await op.process(exps) + self.assertEqual(len(kept), 4) + self.assertEqual(metrics["filtered_count"], 0) + self.assertEqual(metrics["skipped_empty_count"], 2) + self.assertEqual(metrics["skipped_short_count"], 2) + self.assertNotIn("repeat_ratio/mean", metrics) + + async def test_threshold_is_strict(self): + # 8 tokens with ngram_size=2 -> 7 bigrams; [1,2,1,2,1,2,3,4] has 4 unique -> 3/7 + tokens = [1, 2, 1, 2, 1, 2, 3, 4] + ratio = compute_repeat_ratio(tokens, 2) + op_equal = RepetitionFilter(ngram_size=2, max_repeat_ratio=ratio, min_response_tokens=0) + kept, _ = await op_equal.process([make_exp(tokens)]) + self.assertEqual(len(kept), 1) + op_below = RepetitionFilter( + ngram_size=2, max_repeat_ratio=ratio - 1e-6, min_response_tokens=0 + ) + kept, _ = await op_below.process([make_exp(tokens)]) + self.assertEqual(len(kept), 0) + + async def test_only_generated_tokens_are_inspected(self): + op = RepetitionFilter(ngram_size=4, max_repeat_ratio=0.5, min_response_tokens=32) + # Repetitive prompt must not affect the decision. + exp = Experience( + eid=EID(task=0), + tokens=torch.tensor(repetitive(64) + diverse(64), dtype=torch.long), + prompt_length=64, + reward=1.0, + ) + kept, metrics = await op.process([exp]) + self.assertEqual(len(kept), 1) + self.assertEqual(metrics["repeat_ratio/max"], 0.0) + + # Multi-turn: repetitive observation tokens are masked out, generated tokens diverse. + response = repetitive(64) + diverse(64) + mask = [0] * 64 + [1] * 64 + kept, metrics = await op.process([make_exp(response, action_mask=mask)]) + self.assertEqual(len(kept), 1) + self.assertEqual(metrics["repeat_ratio/max"], 0.0) + + # Multi-turn: diverse observation tokens, repetitive generated tokens -> dropped. + response = diverse(64) + repetitive(64) + kept, metrics = await op.process([make_exp(response, action_mask=mask)]) + self.assertEqual(len(kept), 0) + self.assertEqual(metrics["filtered_count"], 1) + + async def test_disable_with_ratio_one(self): + op = RepetitionFilter(max_repeat_ratio=1.0, min_response_tokens=0) + kept, metrics = await op.process([make_exp(repetitive(64))]) + self.assertEqual(len(kept), 1) + self.assertEqual(metrics["filtered_count"], 0) + + def test_invalid_args(self): + with self.assertRaises(ValueError): + RepetitionFilter(ngram_size=0) + with self.assertRaises(ValueError): + RepetitionFilter(max_repeat_ratio=1.5) + with self.assertRaises(ValueError): + RepetitionFilter(min_response_tokens=-1) + + async def test_create_from_registry(self): + ops = create_operators( + [ + OperatorConfig( + name="repetition_filter", + args={"ngram_size": 3, "max_repeat_ratio": 0.6, "min_response_tokens": 16}, + ) + ] + ) + self.assertEqual(len(ops), 1) + self.assertIsInstance(ops[0], RepetitionFilter) + self.assertEqual(ops[0].ngram_size, 3) + self.assertEqual(ops[0].max_repeat_ratio, 0.6) + self.assertEqual(ops[0].min_response_tokens, 16) + kept, metrics = await ops[0].process([make_exp(diverse(32)), make_exp(repetitive(32))]) + self.assertEqual(len(kept), 1) + self.assertEqual(metrics["filtered_count"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/trinity/buffer/operators/__init__.py b/trinity/buffer/operators/__init__.py index 35baa7dad21..aa00736b1cd 100644 --- a/trinity/buffer/operators/__init__.py +++ b/trinity/buffer/operators/__init__.py @@ -16,6 +16,7 @@ "pass_rate_calculator": "trinity.buffer.operators.mappers.pass_rate_calculator.PassRateCalculator", "data_juicer": "trinity.buffer.operators.data_juicer_operator.DataJuicerOperator", "invalid_reward_filter": "trinity.buffer.operators.filters.reward_filter.InvalidRewardFilter", + "repetition_filter": "trinity.buffer.operators.filters.repetition_filter.RepetitionFilter", }, ) diff --git a/trinity/buffer/operators/filters/repetition_filter.py b/trinity/buffer/operators/filters/repetition_filter.py new file mode 100644 index 00000000000..6a97a3e8f75 --- /dev/null +++ b/trinity/buffer/operators/filters/repetition_filter.py @@ -0,0 +1,121 @@ +"""Filter out experiences whose generated response is highly repetitive.""" + +from typing import Dict, List, Sequence, Tuple + +from trinity.buffer.operators import ExperienceOperatorV1 +from trinity.common.experience import Experience + + +def compute_repeat_ratio(tokens: Sequence[int], ngram_size: int) -> float: + """Compute the duplicate n-gram ratio of a token sequence. + + The ratio is defined as ``1 - unique_ngrams / total_ngrams`` over all sliding + windows of ``ngram_size`` tokens. It is ``0.0`` when every n-gram is distinct and + approaches ``1.0`` when the sequence keeps repeating the same short pattern. + + Args: + tokens: Token ids of the generated response. + ngram_size: Size of the sliding window. + + Returns: + float: Duplicate n-gram ratio in ``[0, 1]``. Sequences shorter than + ``ngram_size`` have no n-grams and score ``0.0``. + """ + total = len(tokens) - ngram_size + 1 + if total <= 0: + return 0.0 + unique = len({tuple(tokens[i : i + ngram_size]) for i in range(total)}) + return 1.0 - unique / total + + +class RepetitionFilter(ExperienceOperatorV1): + """ + Drop experiences whose response has collapsed into repetition. + + Repetition collapse is a common failure mode of on-policy rollouts: the policy keeps + emitting the same pattern instead of terminating. Such trajectories carry almost no + learning signal, so this operator removes them before they enter the training buffer. + + Only the model-generated part of the experience is inspected: tokens before + ``prompt_length`` are ignored, and when ``action_mask`` is available (multi-turn + experiences) tokens with a zero mask (e.g. environment observations) are ignored too. + An experience is dropped when ``duplicate_ngrams / total_ngrams > max_repeat_ratio``. + Empty responses and responses shorter than ``min_response_tokens`` are always kept, + since the ratio is not meaningful on very short sequences. + """ + + def __init__( + self, + ngram_size: int = 4, + max_repeat_ratio: float = 0.5, + min_response_tokens: int = 32, + ) -> None: + """Initialize the repetition filter. + + Args: + ngram_size: Sliding window size (in tokens) used to measure repetition. + max_repeat_ratio: Experiences whose duplicate n-gram ratio is strictly greater + than this value are dropped. ``0.0`` drops any response containing a + repeated n-gram, ``1.0`` disables filtering. + min_response_tokens: Responses with fewer generated tokens than this are kept + without being scored. + """ + if ngram_size < 1: + raise ValueError(f"ngram_size must be >= 1, got {ngram_size}") + if not 0.0 <= max_repeat_ratio <= 1.0: + raise ValueError(f"max_repeat_ratio must be in [0, 1], got {max_repeat_ratio}") + if min_response_tokens < 0: + raise ValueError(f"min_response_tokens must be >= 0, got {min_response_tokens}") + self.ngram_size = ngram_size + self.max_repeat_ratio = max_repeat_ratio + self.min_response_tokens = min_response_tokens + + @staticmethod + def _response_tokens(exp: Experience) -> List[int]: + """Extract the model-generated token ids of an experience.""" + if exp.tokens is None: + return [] + response = exp.tokens[exp.prompt_length :] + if exp.action_mask is not None: + mask = exp.action_mask.bool() + length = min(len(response), len(mask)) + response = response[:length][mask[:length]] + return response.tolist() + + async def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]: + """Drop experiences whose response exceeds the repetition threshold. + + Args: + exps: Experiences to process. + + Returns: + Tuple[List[Experience], Dict]: Kept experiences and filtering metrics. + """ + kept = [] + ratios = [] + skipped_empty = 0 + skipped_short = 0 + for exp in exps: + tokens = self._response_tokens(exp) + if len(tokens) == 0: + skipped_empty += 1 + kept.append(exp) + continue + if len(tokens) < self.min_response_tokens: + skipped_short += 1 + kept.append(exp) + continue + ratio = compute_repeat_ratio(tokens, self.ngram_size) + ratios.append(ratio) + if ratio > self.max_repeat_ratio: + continue + kept.append(exp) + metrics: Dict[str, float] = { + "filtered_count": len(exps) - len(kept), + "skipped_empty_count": skipped_empty, + "skipped_short_count": skipped_short, + } + if ratios: + metrics["repeat_ratio/mean"] = sum(ratios) / len(ratios) + metrics["repeat_ratio/max"] = max(ratios) + return kept, metrics