dockyard_rl.environments.math_environment

Module Contents

Classes

API

class dockyard_rl.environments.math_environment.MathEnvConfig

Bases: typing.TypedDict

num_workers: int

None

stop_strings: NotRequired[list[str] | None]

None

verifier_type: NotRequired[str | None]

None

math_verify_impl: NotRequired[str | None]

None

class dockyard_rl.environments.math_environment.HFVerifyWorker

Initialization

verify(pred_responses: list[str], ground_truths: list[str], return_extracted_answer: bool = False, **kwargs) Union[list[float], tuple[list[float], list[str | None]]]

Verify the correctness of the predicted responses against the ground truth.

Args: pred_responses: list[str]. The predicted responses from the LLM. ground_truths: list[str]. The ground truth responses.

Returns: Union[list[float], tuple[list[float], list[str | None]]]. If return_extracted_answer is False, returns only the scores. If return_extracted_answer is True, returns (scores, extracted_answers).

class dockyard_rl.environments.math_environment.MultilingualMultichoiceVerifyWorker
verify(pred_responses: list[str], ground_truths: list[str], return_extracted_answer: bool = False, **kwargs) Union[list[float], tuple[list[float], list[str | None]]]

Verify the correctness of the predicted responses against the ground truth.

Args: pred_responses: list[str]. The predicted responses from the LLM. ground_truths: list[str]. The ground truth responses.

Returns: Union[list[float], tuple[list[float], list[str | None]]]. If return_extracted_answer is False, returns only the scores. If return_extracted_answer is True, returns (scores, extracted_answers).

class dockyard_rl.environments.math_environment.EnglishMultichoiceVerifyWorker
verify(pred_responses: list[str], ground_truths: list[str], return_extracted_answer: bool = False, **kwargs) Union[list[float], tuple[list[float], list[str | None]]]

Verify the correctness of the predicted responses against the ground truth.

Args: pred_responses: list[str]. The predicted responses from the LLM. ground_truths: list[str]. The ground truth responses.

Returns: Union[list[float], tuple[list[float], list[str | None]]]. If return_extracted_answer is False, returns only the scores. If return_extracted_answer is True, returns (scores, extracted_answers).

class dockyard_rl.environments.math_environment.HFMultiRewardVerifyWorker

Initialization

REWARD_NAMES: list[str]

[‘reward/correctness’, ‘reward/integer’, ‘reward/format’]

verify(pred_responses: list[str], ground_truths: list[str], return_extracted_answer: bool = False, **kwargs) Union[dict[str, list[float]], tuple[dict[str, list[float]], list[str | None]]]

Verify the correctness of the predicted responses against the ground truth.

Args: pred_responses: list[str]. The predicted responses from the LLM. ground_truths: list[str]. The ground truth responses.

Returns: If return_extracted_answer is False, returns a dict mapping reward component names to per-sample scores. If return_extracted_answer is True, returns (scores_dict, extracted_answers).

class dockyard_rl.environments.math_environment.MathEnvironmentMetadata

Bases: typing.TypedDict

ground_truth: str

None

extracted_answer: str | None

None

class dockyard_rl.environments.math_environment.BaseMathEnvironment(cfg: dockyard_rl.environments.math_environment.MathEnvConfig)

Bases: dockyard_rl.environments.interfaces.EnvironmentInterface[dockyard_rl.environments.math_environment.MathEnvironmentMetadata]

WORKER_CLASS_DICT: dict[str, type]

None

shutdown() None
global_post_process_and_metrics(batch: dockyard_rl.distributed.batched_data_dict.BatchedDataDict[Any]) tuple[dockyard_rl.distributed.batched_data_dict.BatchedDataDict[Any], dict[str, float | int]]

Computes metrics for this environment given a global rollout batch.

Every rank will run this function, so you’re free to use distributed calculations if you’d prefer for heavy metrics.

class dockyard_rl.environments.math_environment.MathEnvironment(cfg: dockyard_rl.environments.math_environment.MathEnvConfig)

Bases: dockyard_rl.environments.math_environment.BaseMathEnvironment

WORKER_CLASS_DICT

None

step(message_log_batch: list[dockyard_rl.data.interfaces.LLMMessageLogType], metadata: list[dockyard_rl.environments.math_environment.MathEnvironmentMetadata], return_extracted_answer: bool = False) dockyard_rl.environments.interfaces.EnvironmentReturn[dockyard_rl.environments.math_environment.MathEnvironmentMetadata]

Runs a step in the math environment.

Args: message_log: list[list[dict[str, str]]]. A batch of OpenAI-API-like message logs that represent interactions with the LLM. metadata: list[MathEnvironmentMetadata]. The grader will use the ‘ground_truth’ key to evaluate correctness. The extracted answer will be stored to caculate cons@k.

Returns: EnvironmentReturn: A tuple containing: - list[dict[str, str]]: Observations/responses batch - list[dict]: Updated metadata - list[str]: Next stop strings for the next turn - Tensor: Rewards tensor - Tensor: Done flags tensor

class dockyard_rl.environments.math_environment.MathMultiRewardEnvironment(cfg: dockyard_rl.environments.math_environment.MathEnvConfig)

Bases: dockyard_rl.environments.math_environment.BaseMathEnvironment

WORKER_CLASS_DICT

None

step(message_log_batch: list[dockyard_rl.data.interfaces.LLMMessageLogType], metadata: list[dockyard_rl.environments.math_environment.MathEnvironmentMetadata], return_extracted_answer: bool = False) dockyard_rl.environments.interfaces.EnvironmentReturn[dockyard_rl.environments.math_environment.MathEnvironmentMetadata]

Runs a step in the math environment.

Args: message_log: list[list[dict[str, str]]]. A batch of OpenAI-API-like message logs that represent interactions with the LLM. metadata: list[MathEnvironmentMetadata]. The grader will use the ‘ground_truth’ key to evaluate correctness. The extracted answer will be stored to caculate cons@k.

Returns: EnvironmentReturn: A tuple containing: - list[dict[str, str]]: Observations/responses batch - list[dict]: Updated metadata - list[str]: Next stop strings for the next turn - Tensor: Rewards tensor - Tensor: Done flags tensor