Implementation:Hpcaitech ColossalAI Reward Utils

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Reinforcement Learning, Reward Modeling, RLHF
Last Updated	2026-02-09 00:00 GMT

Overview

Reward processing utilities for validating response structure and extracting answers from model outputs.

Description

This module provides utility functions for the reward computation pipeline in ColossalChat's RLHF system. validate_response_structure checks whether a model's response follows the expected XML-style tag format (e.g., <think>...</think><answer>...</answer>) by verifying tag counts and ordering. extract_solution extracts the final answer from within <answer> tags. extract_boxed_solution extracts content from LaTeX-style \boxed{} expressions, handling nested braces. These functions are adapted from the Logic-RL project and are used to parse structured model outputs for verifiable reward computation.

Usage

Use these utilities in reward functions that need to validate the format of chain-of-thought responses and extract final answers for comparison against ground truth, particularly in math reasoning and logic tasks.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalChat/coati/distributed/reward/reward_utils.py
• Lines: 1-124

Signature

def validate_response_structure(processed_str: str, tags: Dict = None) -> bool:

def extract_solution(solution_str: str) -> Tuple[Optional[str], str]:

def extract_boxed_solution(text: str) -> Optional[str]:

Import

from coati.distributed.reward.reward_utils import validate_response_structure, extract_solution, extract_boxed_solution

I/O Contract

Inputs (validate_response_structure)

Name	Type	Required	Description
processed_str	str	Yes	Processed response string from the model
tags	Dict	No	Dictionary defining expected tags with text and num_occur keys; defaults to think/answer tags

Outputs (validate_response_structure)

Name	Type	Description
return	bool	Whether all formatting requirements are met

Inputs (extract_solution)

Name	Type	Required	Description
solution_str	str	Yes	Raw response string from the language model

Outputs (extract_solution)

Name	Type	Description
return	Tuple[Optional[str], str]	Tuple of (extracted_answer, processed_string); answer is None if no <answer> tags found

Inputs (extract_boxed_solution)

Name	Type	Required	Description
text	str	Yes	A string potentially containing LaTeX-style \boxed{} expressions

Outputs (extract_boxed_solution)

Name	Type	Description
return	Optional[str]	Content inside the last \boxed{}, or None if not found or unmatched

Usage Examples

from coati.distributed.reward.reward_utils import validate_response_structure, extract_solution, extract_boxed_solution

# Validate response structure
response = "<think>Let me work through this...</think><answer>42</answer>"
is_valid = validate_response_structure(response)  # True

# Extract solution from answer tags
answer, processed = extract_solution(response)  # answer = "42"

# Extract from LaTeX boxed notation
boxed_answer = extract_boxed_solution("The answer is \\boxed{42}.")  # "42"

Related Pages

• Environment:Hpcaitech_ColossalAI_CUDA_GPU_Environment