# Copyright 2024 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import importlib import importlib.util import os import sys from collections import defaultdict from functools import partial from typing import Callable, Optional, Tuple, TypedDict import torch from transformers import PreTrainedTokenizer from ...protocol import DataProto from .config import RewardConfig class RewardInput(TypedDict): prompt: Optional[str] response: str response_length: int ground_truth: str data_type: Optional[str] problem_type: Optional[str] problem: Optional[str] problem_id: Optional[int] class RewardScore(TypedDict): overall: float format: Optional[float] accuracy: Optional[float] def _float_env(name: str, default: Optional[float]) -> Optional[float]: raw = os.getenv(name) if raw is None or raw == "": return default try: return float(raw) except Exception: return default def _apply_oracle_reward_control( score: RewardScore, is_oracle_row: bool, reward_metrics: defaultdict, ) -> RewardScore: """Optionally reduce the scalar reward of oracle rows only. This lives in the reward manager rather than in a task reward file because only the manager sees the `is_oracle_row` flag. Comparing response against ground_truth inside a task reward would also catch the rare on-policy exact match, which must keep its full reward. Env knobs (disabled by default): ORARL_ORACLE_REWARD_SCALE multiply the oracle overall by this value ORARL_ORACLE_REWARD_CAP cap the oracle overall after scaling Example for tracking: an exact annotation normally scores overall=1.5. ORARL_ORACLE_REWARD_CAP=1.25 keeps it a positive anchor while preventing it from dominating the group standard deviation and the selection step under aggressive reward shaping. """ if not is_oracle_row: return score if not isinstance(score, dict) or "overall" not in score: return score scale = _float_env("ORARL_ORACLE_REWARD_SCALE", 1.0) cap = _float_env("ORARL_ORACLE_REWARD_CAP", None) if scale is None: scale = 1.0 old_overall = float(score.get("overall", 0.0)) new_overall = old_overall * float(scale) if cap is not None: new_overall = min(new_overall, float(cap)) if new_overall != old_overall: score = dict(score) score["overall"] = float(new_overall) reward_metrics["oracle_overall_before_control"].append(old_overall) reward_metrics["oracle_overall_after_control"].append(float(new_overall)) reward_metrics["oracle_reward_scale"].append(float(scale)) if cap is not None: reward_metrics["oracle_reward_cap"].append(float(cap)) return score SequentialRewardFunction = Callable[[RewardInput], RewardScore] BatchRewardFunction = Callable[[list[RewardInput]], list[RewardScore]] def _build_reward_input( data: DataProto, response_str: str, response_length: int, index: int, prompt_str: Optional[str] = None, ) -> RewardInput: response_prefix = os.getenv("MULTITASK_REWARD_RESPONSE_PREFIX", "") if response_prefix and not response_str.startswith(response_prefix): response_str = response_prefix + response_str non_tensor = data.non_tensor_batch data_type = non_tensor["data_type"][index] if "data_type" in non_tensor else None problem_type = non_tensor["problem_type"][index] if "problem_type" in non_tensor else None problem = non_tensor["problem_reserved_text"][index] if "problem_reserved_text" in non_tensor else None problem_id = non_tensor["problem_id"][index] if "problem_id" in non_tensor else None item = { "prompt": prompt_str, "response": response_str, "response_length": response_length, "ground_truth": non_tensor["ground_truth"][index], "data_type": data_type, "problem_type": problem_type, "problem": problem, "problem_id": problem_id, } # Task rewards may need side-channel metadata that is not part of the text # prompt. Mask-aware segmentation reward, for instance, needs the annotated # mask payload and video metadata to score the predicted prompts. for key in ( "segmentation_output", "fps", "video_second", "resolution", "path", "data_source", "meta", ): if key in non_tensor: item[key] = non_tensor[key][index] return item class SequentialFunctionRewardManagerMixin: reward_fn: SequentialRewardFunction def compute_reward_sequential(self, data: DataProto) -> Tuple[torch.Tensor, dict[str, list[float]]]: reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) reward_metrics = defaultdict(list) prompt_ids = data.batch.get("prompts", None) response_ids = data.batch["responses"] response_length = torch.sum(data.batch["response_mask"], dim=-1) for i in range(len(data)): prompt_str = None if prompt_ids is not None: prompt_str = self.tokenizer.decode(prompt_ids[i], skip_special_tokens=self.config.skip_special_tokens) cur_response_length = int(response_length[i].item()) # avoid tensor indexing error valid_response_ids = response_ids[i][:cur_response_length] response_str = self.tokenizer.decode( valid_response_ids, skip_special_tokens=self.config.skip_special_tokens ) score = self.reward_fn(_build_reward_input(data, response_str, cur_response_length, i, prompt_str)) is_oracle = bool(data.non_tensor_batch.get("is_oracle_row", [False] * len(data))[i]) score = _apply_oracle_reward_control(score, is_oracle, reward_metrics) reward_tensor[i, cur_response_length - 1] = score["overall"] for key, value in score.items(): reward_metrics[key].append(value) return reward_tensor, reward_metrics class BatchFunctionRewardManagerMixin: reward_fn: BatchRewardFunction def compute_reward_batch(self, data: DataProto) -> Tuple[torch.Tensor, dict[str, list[float]]]: reward_inputs = [] prompt_ids = data.batch.get("prompts", None) response_ids = data.batch["responses"] response_length = torch.sum(data.batch["response_mask"], dim=-1) for i in range(len(data)): prompt_str = None if prompt_ids is not None: prompt_str = self.tokenizer.decode(prompt_ids[i], skip_special_tokens=self.config.skip_special_tokens) cur_response_length = int(response_length[i].item()) # avoid tensor indexing error valid_response_ids = response_ids[i][:cur_response_length] response_str = self.tokenizer.decode( valid_response_ids, skip_special_tokens=self.config.skip_special_tokens ) reward_inputs.append(_build_reward_input(data, response_str, cur_response_length, i, prompt_str)) scores = self.reward_fn(reward_inputs) reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) reward_metrics = defaultdict(list) is_oracle_rows = data.non_tensor_batch.get("is_oracle_row", [False] * len(data)) for i, score in enumerate(scores): cur_response_length = int(response_length[i].item()) # avoid tensor indexing error score = _apply_oracle_reward_control(score, bool(is_oracle_rows[i]), reward_metrics) reward_tensor[i, cur_response_length - 1] = score["overall"] for key, value in score.items(): reward_metrics[key].append(value) return reward_tensor, reward_metrics class AutoRewardManager(BatchFunctionRewardManagerMixin, SequentialFunctionRewardManagerMixin): """Reward manager for rule-based reward.""" def __init__(self, config: RewardConfig, tokenizer: PreTrainedTokenizer): if config.reward_function is None: raise ValueError("Reward function is not provided.") if config.reward_function_is_module: try: module = importlib.import_module(config.reward_function) except Exception as e: raise RuntimeError( f"Failed to import reward module {config.reward_function!r}: {e}" ) from e else: if not os.path.exists(config.reward_function): raise FileNotFoundError( f"Reward function file {config.reward_function} not found." ) spec = importlib.util.spec_from_file_location( "custom_reward_fn", config.reward_function, ) if spec is None or spec.loader is None: raise RuntimeError( f"Could not create an import spec for {config.reward_function!r}." ) module = importlib.util.module_from_spec(spec) try: sys.modules["custom_reward_fn"] = module spec.loader.exec_module(module) except Exception as e: raise RuntimeError(f"Failed to load reward function: {e}") from e if not config.reward_function_name or not hasattr( module, config.reward_function_name, ): raise AttributeError(f"Module {module} does not have function {config.reward_function_name}.") reward_fn = getattr(module, config.reward_function_name) reward_name = getattr(module, "REWARD_NAME", "unknown") reward_type = getattr(module, "REWARD_TYPE", "batch") print(f"Using reward function `{config.reward_function_name}` from `{config.reward_function}`.") print(f"Reward name: {reward_name}, reward type: {reward_type}.") self.reward_fn = partial(reward_fn, **config.reward_function_kwargs) self.reward_type = reward_type self.config = config self.tokenizer = tokenizer def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, dict[str, list[float]]]: """Compute reward for a batch of data.""" if self.reward_type == "batch": return self.compute_reward_batch(data) elif self.reward_type == "sequential": return self.compute_reward_sequential(data) else: raise ValueError(f"Unsupported reward type: {self.reward_type}.")