---

name: reward-modeling description: "Work with RewardModel and ImplicitPRM workflows in palm-rlhf-pytorch, including masks, labels, and tiny reward smoke tests." disable-model-invocation: true metadata: disco-role: operating license: MIT

Reward Modeling

Use this sub-skill when a task is about scoring sequences, training rewards, or deriving implicit process rewards from a PaLM backbone.

Route Here For

• Building or inspecting RewardModel for scalar or binned rewards.
• Using prompt_mask or prompt_lengths correctly, with the reminder that they are mutually exclusive.
• Training a reward model with labels and checking the inference path that returns either logits or sampled bin ids.
• Using ImplicitPRM for dense token-level rewards that compare a model against a frozen reference model.
• Running the tiny smoke checks in scripts/tiny_reward_smoke.py and scripts/tiny_implicit_prm_smoke.py.

Do Not Use This For

• Base transformer sizing, generation, or LoRA scope mechanics that belong to palm-modeling.
• PPO, GRPO, TPO, or FlowRL trainer selection and training loops that belong to policy-optimization.
• Full-scale reward datasets or benchmark training. This sub-skill covers the package mechanics and safe smoke checks, not long runs.

Start Here

1. Read references/api-reference.md for the verified constructor and forward signatures.
2. Read references/workflows.md for the scalar reward and implicit process reward recipes.
3. Read references/troubleshooting.md for mask conflicts, binned-vs-logit inference, and shape issues.
4. Run the tiny helper that matches your need:

python sub-skills/reward-modeling/scripts/tiny_reward_smoke.py --device auto
python sub-skills/reward-modeling/scripts/tiny_implicit_prm_smoke.py --device auto

If you are already in this sub-skill directory, use the local scripts/ path instead.

Operating Notes

• RewardModel deep-copies the supplied PaLM and can add a named reward LoRA scope.
• For inference, remember that a binned reward model samples bin ids by default unless sample_from_bins=False is passed.
• Use labels shaped to the pooled reward output: floats for scalar MSE and integer class ids for binned cross entropy.
• ImplicitPRM predicts token-level rewards from the log-prob gap between a trainable model and a frozen reference model.
• Prompt and response embeddings are injected before the copied backbone; the pooling mask is separate from the prompt indicator.

Common Prompt Patterns

• "Train a scalar reward model" → use RewardModel(..., num_binned_output=0).
• "Need reward logits instead of sampled bins" → set sample_from_bins=False.
• "Use prompt lengths or a mask" → pick exactly one prompt indicator style and keep the labels compatible.
• "Need step-by-step dense rewards" → use ImplicitPRM and expect the shifted output length.
• "Why is the output one token shorter?" → because the process reward compares source and target tokens after shifting.

Quick Decision Points

• If the task is about a single score per sequence, start with scalar RewardModel.
• If the task is about bins or ratings, use binned RewardModel and choose whether you want logits or sampled ids.
• If the task is about dense process signals, use ImplicitPRM.
• If the task is about prompt/response semantics, inspect the prompt mask and pooled mask separately.
• If the task is about a downstream trainer, route there after the reward model is settled.

Tiny Verification Checklist

A correct tiny check should:

• construct a tiny PaLM backbone;
• run scalar reward training with labels and a prompt indicator;
• run binned reward training or logits inference;
• confirm whether binned inference returns logits or sampled class ids;
• run an ImplicitPRM backward pass;
• confirm the dense reward output shape is (batch, seq_len - 1).

Related Files

• references/api-reference.md for verified signatures and output shapes.
• references/workflows.md for scalar, binned, and implicit-process workflows.
• references/troubleshooting.md for prompt-mask, shape, and inference issues.
• scripts/tiny_reward_smoke.py and scripts/tiny_implicit_prm_smoke.py for safe runnable smoke checks.