AWS sagemaker documentation change
Summary
Added a new section on reward function design for reinforcement learning, including guidelines, design process, and a search agent example
Security assessment
The change adds general machine learning content about reward function design in RL. There is no mention of security vulnerabilities, fixes, or security features. The content focuses on improving model performance and preventing reward hacking, which is a general ML concept not specific to security.
Diff
diff --git a/sagemaker/latest/dg/model-customize-mtrl-assets.md b/sagemaker/latest/dg/model-customize-mtrl-assets.md index 370b936bc..8453a5b3c 100644 --- a//sagemaker/latest/dg/model-customize-mtrl-assets.md +++ b//sagemaker/latest/dg/model-customize-mtrl-assets.md @@ -7 +7 @@ -Prompt dataset formatExample 1: Simple Q&A Dataset (Plain Text)Example 2: Search/Reasoning with Tool UseExample 3: SQL Generation (Multi-Turn with Complex Context)Best Practices +Prompt dataset formatExample 1: Simple Q&A Dataset (Plain Text)Example 2: Search/Reasoning with Tool UseExample 3: SQL Generation (Multi-Turn with Complex Context)Best PracticesReward function design @@ -246,0 +247,105 @@ Minimum examples at least equal to `training_batch_size`. 10x+ your batch size f +## Reward function design + +Reward function design is critical for providing effective learning signals in complex, multi-step agent systems. When designing reward functions for multi-turn RL, consider the following guidelines. + + * **Start with outcome-based rewards.** Score the final result first to establish a clean and reliable baseline before adding intermediate rewards or reward shaping. + + * **Consider continuous rewards over binary rewards.** Continuous rewards can provide clearer partial-credit signals, but are easy to game. Binary rewards are preferred when partial credit is hard to define or when a clean baseline is needed. + + * **Use shaping rewards carefully.** Shaping rewards can guide learning, but they should be used sparingly because overly strong or misaligned shaping may teach shortcuts. + + * **Guard against reward hacking.** Make rewards difficult to exploit, and verify that the model is solving the real task rather than gaming the scoring rule. + + * **Validate before training.** Test the reward function on real trajectories before training to catch bugs, loopholes, or misleading signals. + + * **Monitor behavioral metrics, not just reward.** Track metrics such as completion rate, turn count, tool use, and overfitting gap to ensure the model is improving in the intended way. + + + + +### Reward design process + + 1. Define what success looks like and determine whether it can be scored automatically. + + 2. Evaluate the base model to establish a baseline success rate. + + 3. Design reward tiers: positive rewards for success, zero rewards for failure, and negative rewards for degenerate behavior. + + 4. Handle edge cases explicitly, including timeouts, environment errors, malformed outputs, and empty responses. + + 5. Check each reward component for potential reward hacking. + + 6. Validate on real trajectories before training. + + 7. Monitor alongside behavioral metrics during training. + + 8. Iterate based on initial results. + + + + +In practice, a reward function takes the full message history of an episode as input and returns two outputs: a scalar reward (a floating-point score measuring trajectory quality, with higher values indicating better performance) and a metrics dictionary for logging, debugging, and monitoring. + +### Example: Search agent reward function + +The following example shows a reward function for an agent that answers questions using search. It demonstrates outcome evaluation, format shaping, and answer correctness checking. + + + class TextAnswerReward: + """Reward function to check text answer against gold answers. + + formula: format_coef * (correct_format - 1) + correct_answer + """ + + gold_answers: list[str] + format_coef: float = 0.1 + + async def __call__(self, history: list[Message]) -> tuple[float, dict[str, float]]: + """Grade the completed episode by checking the final assistant message.""" + final_message = None + for msg in reversed(history): + if msg.get("role") == "assistant": + final_message = msg + break + + if final_message is None: + return 0.0, {"format": 0.0, "correct": 0.0} + + content = get_text_content(final_message) + + correct_format = float(self._extract_answer(content) is not None) + correct_answer = float(self._check_answer(content)) + + reward = self.format_coef * (correct_format - 1) + correct_answer + return reward, {"format": correct_format, "correct": correct_answer} + + def _extract_answer(self, text: str) -> str | None: + if "Answer:" not in text: + return None + parts = text.split("Answer:") + if len(parts) != 2: + return None + return parts[1].strip() + + def _check_answer(self, text: str) -> bool: + model_answer = self._extract_answer(text) + if model_answer is None or len(self.gold_answers) == 0: + return False + for gold in self.gold_answers: + if normalize_answer(model_answer) == normalize_answer(gold): + return True + return False + +This reward function includes the following key design choices: + + * **Correctness dominates.** A correct answer always scores higher than an incorrect one, regardless of format. + + * **Format is a small shaping signal.** The format coefficient (0.1) is 10% of the outcome reward, small enough that the model cannot profit from format compliance alone, but large enough to steer it toward parseable outputs. + + * **Wrong format with wrong answer is mildly penalized.** The -0.1 score creates a small gradient away from completely unstructured outputs, without overwhelming the learning signal. + + * **No answer is treated as incorrect with bad format.** If the model never produces an assistant message, the function returns 0.0, distinguishing it from the active penalty of -0.1 for a present but malformed response. + + + +