---

name: policy-optimization description: "Guide PPO, GRPO, TPO, and FlowRL post-training loops for palm-rlhf-pytorch, including bounded RLHF smoke checks." disable-model-invocation: true metadata: disco-role: operating license: MIT

Policy Optimization

Use this sub-skill when a task is about post-training a PaLM policy with the package trainers: PPO RLHFTrainer, ActorCritic, GRPO, TPO, or FlowRL.

Route Here For

• Default PPO RLHF workflows that use the root exports RLHFTrainer and ActorCritic.
• Trainer selection questions: PPO vs GRPO vs TPO vs FlowRL.
• Bounded tiny smoke checks that verify trainer construction, prompt token id handling, and a single short training or generation step.
• Side-effect awareness for replay buffers, partition functions, and Accelerate-wrapped modules.

Do Not Use This For

• Base PaLM construction, generation, or LoRA scope wiring; route those to palm-modeling.
• Reward training or implicit process reward workflows; route those to reward-modeling.
• Long training, benchmarks, or full RLHF scaling runs. This sub-skill covers mechanics and smoke checks, not expensive experiments.

Start Here

1. For API facts, read references/api-reference.md.
2. For trainer choice, read references/trainer-selection.md.
3. For bounded recipes and the tiny smoke workflow, read references/workflows.md.
4. For known failures and source footguns, read references/troubleshooting.md.
5. To verify the installed package, run the bundled helper instead of the source example script:

python sub-skills/policy-optimization/scripts/tiny_rlhf_smoke.py --device auto --train-smoke

If you only want a construction check first, add --construct-only.

Operating Notes

• The root export from palm_rlhf_pytorch import RLHFTrainer, ActorCritic is the PPO path.
• GRPO, TPO, and FlowRL are module-qualified trainers and should be selected intentionally.
• Every trainer expects exactly one prompt input style: prompts, prompts_path, or prompt_token_ids.
• Raw text prompts require a tokenizer that returns padded token-id tensors.
• Keep smoke runs tiny. The source defaults are much larger than a verification check needs.
• examples.py is useful evidence, but it is not the preferred runtime path for this sub-skill because it depends on an undeclared lion_pytorch extra and uses expensive defaults.
• PPO updates actor and critic/value parameters; the newer trainer variants change the training state and storage pattern in different ways.

Common Prompt Patterns

• "Which trainer should I use?" → start with trainer-selection.md.
• "Need a tiny RLHF smoke test" → use scripts/tiny_rlhf_smoke.py with --construct-only or --train-smoke.
• "Using raw prompts" → provide a tokenizer and confirm padding.
• "Need prompt-token-id wiring" → use the root PPO path first because it matches the README example shape.
• "Need a critic-free variant" → compare GRPO with TPO and FlowRL before choosing.

Quick Decision Points

• If the request says "RLHF" without extra qualifiers, use the root PPO path first.
• If the request says "group-relative" or "critic-free", compare GRPO next.
• If the request says "target policy optimization" or mentions replay buffers, compare TPO.
• If the request says "flow balance" or "reward distribution matching", compare FlowRL.
• If the request only needs importability or a minimal construction check, use the bundled smoke helper with --construct-only.

Tiny Verification Checklist

A correct tiny check should:

• build a tiny PaLM and scalar RewardModel;
• create prompt_token_ids as a padded integer tensor;
• instantiate the PPO trainer with tiny hyperparameters;
• optionally run one bounded update and one generate call;
• keep max_norm=None for the smoke path;
• avoid the source example's large defaults and extra dependency.

Related Files

• references/api-reference.md for constructor signatures and method shapes.
• references/trainer-selection.md for PPO vs GRPO vs TPO vs FlowRL choice guidance.
• references/workflows.md for bounded recipes and smoke-script usage.
• references/troubleshooting.md for prompt assertions, replay-buffer side effects, and known source footguns.
• scripts/tiny_rlhf_smoke.py for the safe runnable PPO smoke check.