Bandit environment formulation for RLHF

Training language models to follow instructions with human feedbackinherited

The RL environment used for PPO fine-tuning, in which a random customer prompt is presented, the policy produces one response, a reward is assigned by the reward model, and the episode ends immediately.

Most reinforcement learning imagines a long game: act, see the world change, act again. InstructGPT's setup is far simpler. The model sees one prompt, writes one response, receives one score, and the episode is over.

That one-step framing gets worked through in general before landing on what it costs the surrounding machinery. It is worth following into generalized advantage estimation, where the same simplicity leaves half the usual apparatus doing nothing.

A one-step decision problem, in general

A bandit is the simplest setting in reinforcement learning: an agent observes a context, takes one action, receives one reward, and the interaction ends, with no state transition and therefore no need for multi-step credit assignment. This contrasts with a full Markov decision process, where an action changes the state and the agent must reason about a whole trajectory of future rewards. Contextual bandits, where the reward-generating action depends on an observed context, predate this paper by decades, tracing to the classical multi-armed bandit problem and its context-dependent extensions in the RL and online-learning literature. The concept is inherited: InstructGPT does not introduce the bandit formulation but adopts it as the environment for its PPO fine-tuning stage (RLHF policy optimization) stage.

How PPO fine-tuning instantiates it

InstructGPT's RL environment "presents a random customer prompt and expects a response to the prompt. Given the prompt and response, it produces a reward determined by the reward model and ends the episode" (instructgpt, §"3.5 Models", p. 9). One prompt, one sampled response, one scalar score from the Reward model (RM), then termination: there is no second turn, no state the policy must track across exchanges, and no long-horizon credit-assignment problem of the kind bandit formulations were built to avoid worrying about. As a child concept of PPO fine-tuning stage (RLHF policy optimization), this is the environment shape that the rest of the PPO machinery, PPO value function, Generalized advantage estimation (GAE), and the Per-token KL penalty from the SFT model, is built to operate inside.

Consequences for the surrounding machinery

The single-step structure quietly deactivates part of the RL toolchain built for longer horizons. As the Bandit environment formulation for RLHF — renders discounting moot in → Generalized advantage estimation (GAE) edge explains, GAE's discount factor exists to down-weight distant future rewards, but with no future beyond the current response, InstructGPT applies no discount when estimating the generalized advantage (instructgpt, §"C.4 Details of RLHF training", p. 42); only GAE's variance-reduction role across tokens remains load-bearing. The bandit structure also explains why the Per-token KL penalty from the SFT model matters as much as it does: since the reward model contributes exactly one number per episode, delivered only once the response is complete, the per-token KL term is the objective's only source of dense, mid-generation feedback, as the Per-token KL penalty from the SFT model — densifies the terminal reward of → Bandit environment formulation for RLHF edge describes.