Building a Toy RLHF System with Markov Processes: A Hands-On Guide

RLHF has become a cornerstone technique for fine-tuning LLMs, enabling them to align with human preferences. At its core, RLHF leverages a Markov Decision Process (MDP) framework, where an agent learns to make sequential decisions based on rewards.

In this article, we’ll walk through a simple implementation of an RLHF system using Python and PyTorch. This toy example models sentence generation as an MDP, training a policy to produce positive-sounding sentences.

Understanding the Concept

Before coding, let’s understand. An MDP consists of states, actions, transitions, and rewards. Here:

states -> current sequence of words

actions -> the next words to add

rewards -> simulate human feedback (e.g., +1 for positive sentences, -1 for negative ones).

The Markov property ensures the next state depends only on the current state and action, mirroring how LLMs generate text autoregressively. Our goal is to train a policy network to maximize cumulative rewards over time using the REINFORCE algorithm, a policy gradient method.

This project is challenging due to variable-length sequences, sparse rewards, and the need to stabilize training in a high-dimensional action space — even with a small vocabulary.

This code runs on a Mac with an M4 chip (or any system with Python and PyTorch)

pip3 install torch torchvision numpy

1. Defining Hyperparameters and Vocabulary

We start by defining the constants that shape our experiment.

# Hyperparameters
VOCAB = ['start', 'the', 'cat', 'dog', 'is', 'happy', 'sad', 'runs', 'success', 'failure', 'end']  # Small vocabulary
VOCAB_SIZE = len(VOCAB)
STATE_DIM = 32  # Embedding dimension for state
HIDDEN_DIM = 64  # Policy network hidden size
MAX_LENGTH = 5  # Max sentence length (episode length)
GAMMA = 0.99  # Discount factor
LR = 0.001  # Learning rate
EPISODES = 1000  # Training episodes
BATCH_SIZE = 32  # For sampling episodes

# Token to index mapping
token_to_idx = {token: idx for idx, token in enumerate(VOCAB)}
idx_to_token = {idx: token for token, idx in token_to_idx.items()}

This section sets up a small vocabulary of 11 words, including start and end tokens. The state dimension (32) and hidden dimension (64) define the policy network’s architecture. MAX_LENGTH limits sentence length, while GAMMA discounts future rewards. The token-to-index mappings enable numerical processing.

2. Designing the Reward Function

The reward function simulates human feedback, a critical component of RLHF.

# Reward function: Simulate human feedback
def compute_reward(sentence):
    # Positive if contains 'happy' or 'success', negative otherwise
    if 'happy' in sentence or 'success' in sentence:
        return 1.0
    else:
        return -1.0

This simple function returns +1 if the sentence contains “happy” or “success” and -1 otherwise. In real RLHF, reward models are more sophisticated, often trained on human-labeled data, but this suffices for our toy example.

3. Building the Policy Network

The policy network predicts the next token’s probability distribution given the current state.

# Policy Network: Simple MLP to predict next token logits from state embedding
class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNetwork, self).__init__()
        self.embedding = nn.Embedding(VOCAB_SIZE, state_dim)  # Embed tokens
        self.fc1 = nn.Linear(state_dim * MAX_LENGTH, hidden_dim)  # Flatten sequence
        self.fc2 = nn.Linear(hidden_dim, action_dim)  # Output logits

    def forward(self, state):
        # State is list of token indices, pad to MAX_LENGTH
        state_padded = state + [0] * (MAX_LENGTH - len(state))  # Pad with 'start' index 0
        state_tensor = torch.tensor(state_padded, dtype=torch.long)
        emb = self.embedding(state_tensor).view(-1)  # Flatten embeddings
        x = F.relu(self.fc1(emb))
        logits = self.fc2(x)
        return logits

This MLP embeds the state (a sequence of token indices) into a fixed-size vector, processes it through a hidden layer, and outputs logits over the vocabulary. Padding ensures consistent input sizes, a practical compromise for variable-length sequences.

4. Sampling Actions and Generating Episodes

We need functions to sample actions and simulate episodes.

# Function to sample action from policy
def sample_action(policy, state):
    logits = policy(state)
    probs = F.softmax(logits, dim=-1)
    action = torch.multinomial(probs, 1).item()
    return action, torch.log(probs[action])  # Action and log prob

# Generate an episode
def generate_episode(policy):
    state = [token_to_idx['start']]  # Initial state
    log_probs = []
    actions = []
    rewards = [0] * MAX_LENGTH  # Rewards are 0 until end

    for t in range(MAX_LENGTH):
        action, log_prob = sample_action(policy, state)
        actions.append(action)
        log_probs.append(log_prob)
        state.append(action)  # Update state

        if idx_to_token[action] == 'end':
            break  # Early termination if 'end' is chosen

    # Compute reward at end
    sentence = [idx_to_token[idx] for idx in state[1:]]  # Exclude start
    reward = compute_reward(sentence)
    rewards[-1] = reward  # Assign to last step

    return state, actions, log_probs, rewards

sample_action uses the policy to select the next token probabilistically, returning the action and its log probability for gradient computation. generate_episode builds a sentence step-by-step, assigning a reward at the end based on the full sequence.

5. Computing Discounted Returns

Rewards are discounted to prioritize immediate gains.

# Compute discounted returns
def compute_returns(rewards):
    returns = []
    G = 0
    for r in reversed(rewards):
        G = r + GAMMA * G
        returns.insert(0, G)
    return returns

6. Training the Policy

The training loop optimizes the policy using REINFORCE.

# Training loop
def train():
    policy = PolicyNetwork(STATE_DIM, HIDDEN_DIM, VOCAB_SIZE)
    optimizer = optim.Adam(policy.parameters(), lr=LR)

    for episode in range(EPISODES):
        state, actions, log_probs, rewards = generate_episode(policy)
        returns = compute_returns(rewards)
        returns = torch.tensor(returns, dtype=torch.float32)

        # Policy gradient loss: -sum(log_prob * return)
        loss = 0
        for log_prob, G in zip(log_probs, returns):
            loss += -log_prob * G

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if episode % 100 == 0:
            sentence = ' '.join([idx_to_token[idx] for idx in state[1:]])
            print(f"Episode {episode}: Sentence: '{sentence}' | Reward: {rewards[-1]}")

    return policy

This loop generates episodes, computes losses based on log probabilities and returns, and updates the policy. It prints progress every 100 episodes.

7. Running the Script

Finally, we execute the training and test the policy.

# Run training
if __name__ == "__main__":
    trained_policy = train()

    # Test: Generate a sample sentence after training
    print("\nGenerating a sample sentence with trained policy:")
    state = [token_to_idx['start']]
    for _ in range(MAX_LENGTH):
        action, _ = sample_action(trained_policy, state)
        state.append(action)
        if idx_to_token[action] == 'end':
            break
    sentence = ' '.join([idx_to_token[idx] for idx in state[1:]])
    print(f"Sample: '{sentence}' | Reward: {compute_reward([idx_to_token[idx] for idx in state[1:]])}")

Running and Results

Save the code as rlhf.py and run it with python3 rlhf.py. On an M4 Mac, you might see output like:

Episode 0: Sentence: 'failure sad runs start sad' | Reward: -1.0
Episode 100: Sentence: 'happy happy is happy is' | Reward: 1.0
Episode 900: Sentence: 'happy happy happy happy happy' | Reward: 1.0
Generating a sample sentence with trained policy:
Sample: 'is happy happy happy happy' | Reward: 1.0

The policy learns to favor “happy”-rich sentences, reflecting the reward function’s influence. This convergence validates the MDP and RLHF setup.

My Conclusion:

This project offers a glimpse into RLHF’s mechanics within an MDP framework, bridging Markov processes and LLM training. Experiment with the code, tweak the reward function, and share your findings!