Reinforcement Learning Notes

RL is an important part of post-training. These notes begin an exploration of its fundamentals.

A Demo of Basic RL Concepts

RL continually improves policies, critics, and related models through actor-environment interactions, actions, and resulting states and rewards.

This example moves a ball left or right from position 0 toward position 2. The code is annotated; several points matter:

• Reward model: Rewards are simple functions of state. To motivate a critic, state 2 has the highest reward while state 1 deliberately has the lowest. An actor considering only immediate rewards would never reach 2.
  • Rewards may be deterministic or nondeterministic. Defining them well, especially intermediate rewards, is crucial.
• Critic model: Critic values capture long-term value rather than only immediate rewards. They compress potential future rewards, so a step’s value is reward + gamma * critic-value.
  • Critic values may be enumerable, as in this example where all possibilities can be visited and updated, or estimated when enumeration is impossible.
• The policy outputs log probabilities with shape (1,T), where T is the action dimension, then introduces randomness similar to LLM temperature.
• Training
  • Rollout: One or more actor steps produce a trajectory containing state, action, logprobs, reward, critic values, and other information.
  • The reward model generally does not participate in RL training.
  • Critic update:
    • Calculate an advantage using adv = step["reward"] + gamma * V[s_next] - V[s].
      • adv > 0 means V[s] was underestimated; adv < 0 means it was overestimated.
      • Train the model using this delta.
  • Policy update:
    • The advantage indicates whether to encourage or suppress the action, changing future log probabilities.
• This minimal demo represents reward, critic, and policy with Python dictionaries. Real training is more complex and varies by method.

1
import numpy as np
2
import random
3

4
# 环境的描述（-2 到 2 共有 5 个点）
5
states = [-2, -1, 0, 1, 2]
6
# 动作集合（左，不动，右）
7
actions = [-1, 0, 1]
8

9
# 执行 action
10
def transition(s, a):
11
    return max(-2, min(2, s + a))
12

13
# reward model 定义
14
def reward(s):
15
    if s == 2:
16
        return 10
17
    if s == 1:
18
        return -5
19
    return -1
20

21
# policy logits: state -> action logits
22
policy_logits = {
23
    s: np.zeros(len(actions)) for s in states
24
}
25

26
def softmax(x):
27
    e = np.exp(x - np.max(x))
28
    return e / e.sum()
29

30
# policy model,
31
def policy(state):
32
    probs = softmax(policy_logits[state])
33
    return probs
34

35
V = {s: 0.0 for s in states}
36
gamma = 0.9
37

38
trajectory = []
39

40
# 一次 rollout
41
state = 0
42
for t in range(100):
43
    probs = policy(state)
44
    action_idx = np.random.choice(len(actions), p=probs)
45
    action = actions[action_idx]
46

47
    logprob = np.log(probs[action_idx] + 1e-8)
48

49
    next_state = transition(state, action)
50
    r = reward(next_state)
51

52
    trajectory.append({
53
        "state": state,
54
        "action": action,
55
        "action_idx": action_idx,
56
        "reward": r,
57
        "logprob": logprob
58
    })
59

60
    state = next_state
61

62

63
# 计算 advantage, 作为 policy model 的输入
64
advantages = []
65
for step in trajectory:
66
    s = step["state"]
67
    a = step["action"]
68
    s_next = transition(s, a)
69
    td_target = step["reward"] + gamma * V[s_next]
70
    advantage = td_target - V[s]
71
    advantages.append(advantage)
72

73
# 利用 advantages 更新 policy model
74
lr = 0.1
75
for step, adv in zip(trajectory, advantages):
76
    s = step["state"]
77
    a_idx = step["action_idx"]
78
    logp = step["logprob"]
79

80
    # policy gradient: ∇ logπ(a|s) * advantage
81
    policy_logits[s][a_idx] += lr * adv
82

83
# 利用 advantages 更新 critic model
84
alpha = 0.1
85
for step, adv in zip(trajectory, advantages):
86
    s = step["state"]
87
    V[s] += alpha * adv
88

89
for i, step in enumerate(trajectory):
90
    print(
91
        f"t={i}, s={step['state']}, a={step['action']}, "
92
        f"r={step['reward']}, logp={step['logprob']:.3f}, "
93
        f"adv={advantages[i]:.3f}"
94
    )

• Beyond the demo, a reference model constrains the policy from drifting. I do not yet understand whether it is necessary: rewards and critics might serve similar purposes, while another model increases complexity and stability risks.
  • The reference may be the initial policy, with KL divergence defining a loss that discourages excessive drift.

PPO vs DPO vs GRPO

PPO, DPO & GRPO: Reinforcement Learning Techniques for Training LLMs is a helpful article.

• PPO (Proximal Policy Optimization): Keeps policy updates close to the pretrained model in this interpretation.
  • Similar overall to the demo, but policy updates use a clipping ratio. For clip=0.1, the change is limited to 10% of log probabilities, or is it model weights? I need to clarify this.
  • Advantage: A(s,a) = Q(s,a) - V(s), where Q(s,a) represents the average reward following transition(s,a).
• DPO (Direct Preference Optimization): Optimizes directly through comparisons.
  • No reward model; the dataset contains user preferences, like choosing between two ChatGPT answers.
  • The policy loss compares the two. If the user prefers A but the model favors B, increase A’s log probabilities and reduce B’s, using the KL-divergence formulation.
  • A reference model supplies a baseline.
• GRPO (Group Relative Policy Optimization): Applies a group-relative preference idea to policy optimization.
  • Rank a group and compare pairs in a DPO-like way. In this description, DPO compares against an off-policy reference, while GRPO compares using its own policy.

Viewed as an evolution:

• PPO is intuitive and closest to traditional RL, but its structure is complex and requires an additional reward model.
• DPO removes the reward model but needs a good human-labeled preference dataset.
• GRPO simplifies further and uses ranking data more fully: four ranked items yield 3+2+1=6 pairs. In this account, a reference-model baseline is no longer needed.

OpenRLHF / Slime / VeRL

• Rigorous OpenRLHF formula derivations: https://zhuanlan.zhihu.com/p/7461863937
• An OpenRLHF source walkthrough at https://github.com/OpenRLHF/OpenRLHF mainly explains grouping inference and training models with Ray.
• The Slime code I have read follows a similar approach to OpenRLHF.

Slime

• https://github.com/THUDM/slime
• Reading Slime seems a good way to learn most RL infrastructure fundamentals. Many details deserve closer inspection.

With the basic RL concepts understood, Slime’s main flow is relatively simple, though it involves many capabilities. Starting from train.py:

1. Build Ray placement groups and reserve resources. Rollout inference and actor training can be colocated or separate, affecting whether weights must be transferred across groups.
2. Initialize TrainGroup. Critic and actor are placed in a process group with ranks 1 and 0 and world_size=2. Earlier ordered resource placement aims to put them on different ranks under the same IP:

1
def connect(self, critic_group):
2
    return ray.get(
3
        [
4
            actor.connect_actor_critic.remote(critic)
5
            for actor, critic in zip(self._actor_handlers, critic_group._actor_handlers, strict=False)
6
        ]
7
    )

3. Initialize actor and critic models.
4. Begin training by generating data with generate_rollout. SGLang generation uses many asynchronous operations for speed; sft_rollout is a simpler implementation reading data_buffer directly. Generate rollouts and rewards with custom reward-model types, returning list[list[Sample]]. Normalize rewards and perform other processing to produce train_data, then apply DDP.

1
while state.remaining_batch_size < target_data_size:
2
    # get samples from the buffer and submit the generation requests.
3
    samples = data_source(args.over_sampling_batch_size)
4
    state.submit_generate_tasks(samples)
5

6
# wait for the generation to finish
7
done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED)

5. Each training actor fetches its data. This seems inefficient: why not push sharding down to the source instead of first putting data in Ray’s memory store and copying it to each shard?
6. Run the critic forward pass, with additional complexity for context parallelism. Process rewards, invoke Megatron’s forward pass, broadcast critic values to the actor in the same process group, and train the critic.

1
for value in values:
2
    handles.append(dist.broadcast(value, src=1, group=group, async_op=True))
3

4
if args.kl_coef != 0 or args.use_kl_loss:
5
    if not log_probs:
6
        log_probs = [torch.empty_like(value) for value in values]
7
    if not ref_log_probs:
8
        ref_log_probs = [torch.empty_like(value) for value in values]
9
    for ref_log_prob, log_prob in zip(ref_log_probs, log_probs, strict=False):
10
        handles.append(dist.broadcast(log_prob, src=0, group=group, async_op=True))
11
        handles.append(dist.broadcast(ref_log_prob, src=0, group=group, async_op=True))

7. The actor follows a similar flow.

Next, some design details.

Router

When initializing RolloutManager for SGLang, a router is created to track SGLang servers and provide health checks. It also exposes three interfaces:

1
def _setup_routes(self):
2
    """Setup all the HTTP routes"""
3

4
    # 在 SGLangRolloutEngine 初始化的时候进行注册(之前提到过 SGLangRolloutEngine 是一组 Ray Actor，独立部署)
5
    self.app.post("/add_worker")(self.add_worker)
6
    self.app.get("/list_workers")(self.list_workers)
7

8
    # 这里 router支持了一个 middleware 能力，在这个方法中主要是利用 radix tree 来做 token cache
9
    self.app.post("/retrieve_from_text")(self.retrieve_from_text)
10

11
    # 所有其他请求直接 route 给对应的 sglang server
12
    self.app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])(self.proxy)

SGLang Rollout

• Rollout generation uses many coroutines and asynchronous operations, managed with GenerateState.
• Call the SGLang server’s generate method directly.
• Generate rewards after rollout.

Speculative Decoding

• Speculative decoding uses a small model to generate tokens when the large model’s decoding is slow, then one large-model forward pass verifies T consecutive predictions.
• I will inspect Slime’s implementation later; it involves Megatron modifications.

FSDP Backend

• In addition to Megatron, PyTorch FSDP is supported.
• During initialization, rank 0 performs dist.broadcast(state_dict, src=0) to prevent other ranks from loading the model again.
• pack_sequencees: Pack a rollout batch to reduce subsequent padding costs.
• Collect metrics through dist.all_gather_object.

weights updater & offload

// TODO

Multi Turn Rollout

Suitable for agent and robotics workloads. // TODO

Quantitization

Quantization commonly reduces GPU-memory and bandwidth demands, but introduces training/inference discrepancies, leading to many possible combinations. // TODO

OpenTinker(RL-as-a-Service)

• https://github.com/open-tinker/OpenTinker