dockyard_rl.algorithms.reward_functions

Reward shaping functions for Project Dockyard.

Currently supports:

  • DAPO-style overlong response penalty (linear decay over buffer zone)

  • Stop-properly penalty (scale reward of truncated responses)

  • Invalid-action penalty (per-sample subtraction from rollout-time invalid-action / malformed-thinking verdict counts, #2656)

Module Contents

Classes

RewardShapingConfig

Configuration for reward function processing.

Functions

apply_reward_shaping

Apply reward shaping penalties to the batch.

apply_invalid_action_penalty

Subtract the per-sample invalid-action / malformed-thinking penalty.

apply_message_span_advantage_penalties

Overwrite the advantage of flagged assistant-message token spans (N2, #2800).

API

class dockyard_rl.algorithms.reward_functions.RewardShapingConfig

Bases: typing.TypedDict

Configuration for reward function processing.

Enables custom reward shaping, currently supporting DAPO-style penalties for responses that exceed the maximum response length threshold.

Initialization

Initialize self. See help(type(self)) for accurate signature.

enabled: bool

None

overlong_buffer_length: NotRequired[int]

None

overlong_buffer_penalty: NotRequired[float]

None

max_response_length: NotRequired[int]

None

stop_properly_penalty_coef: NotRequired[float | None]

None

dockyard_rl.algorithms.reward_functions.apply_reward_shaping(batch: dockyard_rl.distributed.batched_data_dict.BatchedDataDict, cfg: dockyard_rl.algorithms.reward_functions.RewardShapingConfig) → dockyard_rl.distributed.batched_data_dict.BatchedDataDict

Apply reward shaping penalties to the batch.

Supports two mutually exclusive modes:

Stop-properly penalty (stop_properly_penalty_coef is set): Scales the reward of truncated responses by the given coefficient. Overlong buffer parameters are ignored and a warning is printed.

DAPO overlong penalty (stop_properly_penalty_coef is None): Linearly penalises responses that exceed max_response_length - overlong_buffer_length tokens. Based on https://arxiv.org/pdf/2503.14476.

Args: batch: BatchedDataDict containing “total_reward” and, for DAPO mode, “message_log”. cfg: RewardShapingConfig.

Returns: The same BatchedDataDict with “total_reward” updated in-place.

dockyard_rl.algorithms.reward_functions.apply_invalid_action_penalty(batch: dockyard_rl.distributed.batched_data_dict.BatchedDataDict, cfg: InvalidActionPenaltyConfig | None, step: Optional[int] = None) → dockyard_rl.distributed.batched_data_dict.BatchedDataDict

Subtract the per-sample invalid-action / malformed-thinking penalty.

Two paths:

  • Graded / routed (when batch carries message_log — the colocated sync / async paths): subtract only the reward-locus penalties resolved from the per-message invalid_action_violations stamped at rollout time, under penalty_mode and the per-type severities / step scale. The advantage-locus violations are applied to the advantage tensor in grpo (see apply_message_span_advantage_penalties), so the two never stack under "auto".

  • Legacy counts (when message_log is absent — the data-plane driver_carry): subtract from the per-sample invalid_action_count / malformed_thinking_count. The data-plane path therefore applies all penalties at the reward level (no span info).

Disabled config is a strict no-op.

Args: batch: BatchedDataDict with “total_reward” and either “message_log” (graded) or the two count fields (legacy). cfg: InvalidActionPenaltyConfig or None. step: Optional training step for the penalty step-scale (graded path).

Returns: The same BatchedDataDict with “total_reward” updated in-place.

dockyard_rl.algorithms.reward_functions.apply_message_span_advantage_penalties(advantages: torch.Tensor, message_logs: list, cfg: InvalidActionPenaltyConfig | None, step: Optional[int] = None) → tuple[torch.Tensor, dict]

Overwrite the advantage of flagged assistant-message token spans (N2, #2800).

For each sample’s message log, a token offset is accumulated over the messages’ token_ids; for every message carrying advantage-locus violations (under cfg["penalty_mode"]), advantages[i, offset:offset+msg_len] is overwritten with the negative summed penalty for those violations (severity * base * step scale). Mirrors the per-message-span credit assignment of NeMo-RL #2800 on dockyard’s typed violations. A disabled config or no flagged span is a no-op.

Returns the (in-place modified) advantages and a metrics dict.