RL

This phase starts with a policy \(\pi\) and calculates its value \(v_{\pi}\).

• Input: The current fixed policy \(\pi_k\).
• Process: We iterate to find the value by performing "Value Updates" repeatedly.
  1. Initialize a value estimate \(v\) (e.g., \(v=0\)).
  2. Apply the Bellman Expectation Operator repeatedly (infinite steps): \[ v_{i+1} = r_{\pi_k} + \gamma P_{\pi_k} v_i \]
  3. Stop when \(v\) converges (i.e., \(i \to \infty\)).
• Output: The converged state-value function \(v_{\pi_k}\).

Once we know exactly how good the current policy \(\pi_k\) is (from Phase 1), we make it greedy.

• Process: \[ \pi_{k+1} = \arg\max_{\pi} \left( r_{\pi} + \gamma P_{\pi} v_{\pi_k} \right) \]

Instead of forcing the environment to start randomly, we force the agent to behave randomly occasionally.

• Goal: Ensure all state-action pairs are visited "sufficiently many times" without external resets.
• Mechanism: Uses a Soft Policy ( \(\epsilon\) -greedy), meaning there is always a non-zero probability of taking any action in any state.

The algorithm balances these two conflicting goals via the parameter \(\epsilon\) (epsilon).

• Exploration ( The \(\epsilon\) component)
  • Purpose: To discover new strategies and ensure the agent does not get stuck in a suboptimal loop. By taking random actions, the agent eventually wanders into every possible state, fulfilling the coverage requirement that "Exploring Starts" used to handle.
  • Mechanism: With probability \(\epsilon\), the agent ignores its knowledge and chooses randomly.
• Exploitation (The \(1-\epsilon\) component)
  • Purpose: To maximize rewards based on current knowledge. The agent selects the "Greedy" action (the one with the highest estimated value \(q(s,a)\)).
  • Mechanism: With probability \(1 - \epsilon\), the agent chooses the best known action.

• Continuous Learning: Even if the agent always starts at the same spot (Start Line), the stochastic nature of the policy (\(\epsilon\)) ensures that over many episodes, it will eventually drift into unvisited states.

When we use Action value instead of policy for TD, all discussed algorithms can be expressed as a stochastic approximation update solving a Bellman equation. The unified update rule (Gradient Descent form) is:

\[q_{t+1}(s_t, a_t) = q_t(s_t, a_t) - \alpha_t(s_t, a_t) [ q_t(s_t, a_t) - \bar{q}_t ]\]

• \(q_t(s_t, a_t)\): Current Estimate
• \(\alpha_t\): Learning rate (Step size)
• \(\bar{q}_t\): The Target (The noisy sample derived from environment interaction)
• The term \([q_t - \bar{q}_t]\) represents the error to minimize.

• Sarsa:
  • Target Expression (\(\bar{q}_t\)): \[\bar{q}_t = r_{t+1} + \gamma q_t(s_{t+1}, a_{t+1})\]
  • Equation Aimed to Solve: Bellman Expectation Equation (BE) for \(q_{\pi}\): \[q_{\pi}(s, a) = \mathbb{E} [R_{t+1} + \gamma q_{\pi}(S_{t+1}, A_{t+1}) \mid S_t = s, A_t = a]\]
  • Note: This is an on-policy method using the action \(a_{t+1}\) actually taken by the current policy.
• Q-learning
  • Target Expression (\(\bar{q}_t\)): \[\bar{q}_t = r_{t+1} + \gamma \max_{a} q_t(s_{t+1}, a)\]
  • Equation Aimed to Solve: Bellman Optimality Equation (BOE) for \(q_*\): \[q_{*}(s, a) = \mathbb{E} [R_{t+1} + \max_{a} q_{*}(S_{t+1}, a) \mid S_t = s, A_t = a]\]
  • Note: This is an off-policy method because it updates towards the best possible action (\(\max\)), regardless of the policy actually followed.

• Expected Sarsa
  • Target Expression (\(\bar{q}_t\)): \[\bar{q}_t = r_{t+1} + \gamma \sum_{a} \pi_t(a|s_{t+1})q_t(s_{t+1}, a)\]
  • Equation Aimed to Solve: Bellman Expectation Equation (BE) for \(q_{\pi}\): \[q_{\pi}(s, a) = \mathbb{E} [R_{t+1} + \gamma \mathbb{E}_{A_{t+1}}[q_{\pi}(S_{t+1}, A_{t+1})] \mid S_t = s, A_t = a]\]
  • Note: It reduces variance by taking the expectation over all possible next actions rather than just sampling one.
• N-step Sarsa
  • Target Expression (\(\bar{q}_t\)): \[\bar{q}_t = r_{t+1} + \gamma r_{t+2} + \dots + \gamma^n q_t(s_{t+n}, a_{t+n})\]
  • Equation Aimed to Solve: Bellman Expectation Equation (BE) for \(q_{\pi}\): \[q_{\pi}(s, a) = \mathbb{E} [R_{t+1} + \gamma R_{t+2} + \dots + \gamma^n q_{\pi}(S_{t+n}, A_{t+n}) \mid S_t = s, A_t = a]\]
  • Note: It balances bias and variance by looking \(n\) steps ahead before bootstrapping.
• Monte Carlo (MC)
  • Target Expression (\(\bar{q}_t\)): \[\bar{q}_t = r_{t+1} + \gamma r_{t+2} + \dots \text{ (Full Return } G_t)\]
  • Equation Aimed to Solve: Bellman Expectation Equation (BE) for \(q_{\pi}\): \[q_{\pi}(s, a) = \mathbb{E} [R_{t+1} + \gamma R_{t+2} + \dots \mid S_t = s, A_t = a]\]
  • Note: Can be viewed as the unified expression where \(\alpha_t(s_t, a_t) = 1\), making \(q_{t+1} = \bar{q}_t\) (direct assignment of the return).

Derived from the Policy Gradient theorem, QAC replaces the Monte Carlo (MC) return with a Temporal Difference (TD) value function approximation. The policy parameters are updated using the gradient of the log-probability scaled by the action-value function: \[\theta_{t+1} = \theta_t + \alpha \nabla_{\theta} \ln \pi(a_t \mid s_t, \theta_t) q_t(s_t, a_t)\]

• Policy Function: \(\pi(a|s, \theta_0)\) with initial parameters \(\theta_0\).
• Value Function: \(q(s, a, w_0)\) with initial parameters \(w_0\).
• Learning Rates: \(\alpha_w, \alpha_{\theta} > 0\).
• Goal: Maximize the expected return \(J(\theta)\).

1. Generate Action: Sample \(a_t \sim \pi(a|s_t, \theta_t)\).
2. Observe Reward: Get \(r_{t+1}\) and next state \(s_{t+1}\).
3. Select Next Action: Sample \(a_{t+1} \sim \pi(a|s_{t+1}, \theta_t)\).
4. Actor (Policy Update): Update the policy parameters in the direction of higher rewards: \[\theta_{t+1} = \theta_t + \alpha_{\theta} \nabla_{\theta} \ln \pi(a_t | s_t, \theta_t) q(s_t, a_t, w_t)\]
5. Critic (Value Update): Update the action-value parameters using the semi-gradient TD error: \[w_{t+1} = w_t + \alpha_w \left[ r_{t+1} + \gamma q(s_{t+1}, a_{t+1}, w_t) - q(s_t, a_t, w_t) \right] \nabla_w q(s_t, a_t, w_t)\]

To improve the stability of the Policy Gradient, we introduce a baseline to the update function. Subtracting a state-value function \(v(s)\) from the return reduces variance without introducing bias. The standard gradient update is modified by subtracting the baseline \(v_t(s_t)\): \[ \theta_{t+1} = \theta_t + \alpha \nabla_{\theta} \ln \pi(a_t \mid s_t, \theta_t) [q_t(s_t, a_t) - v_t(s_t)] \]

We define the Advantage Function as the difference between the action-value and the state-value. However, since we often do not know \(q_t\) explicitly, we approximate it using the TD Error \(\delta_t\):

\[ \delta_t(s_t, a_t) = q_t(s_t, a_t) - v_t(s_t) \approx r_{t+1} + \gamma v_t(s_{t+1}) - v_t(s_t) \]

Substituting this back into our update rule gives us a more robust update: \[ \theta_{t+1} = \theta_t + \alpha \nabla_{\theta} \ln \pi(a_t \mid s_t, \theta_t) \delta_t(s_t, a_t) \]

Why use the Advantage Function? The advantage \(\delta_t\) effectively measures how much better an action \(a_t\) was compared to the "average" value of the state. This helps balance exploration and exploitation more effectively than raw returns, as the agent explicitly learns which actions outperform the expected baseline.

• Policy Function (Actor): \(\pi(a|s, \theta_0)\) with initial parameters \(\theta_0\).
• Value Function (Critic): \(v(s, w_0)\) with initial parameters \(w_0\).
• Learning Rates: \(\alpha_w, \alpha_{\theta} > 0\).
• Goal: Learn an optimal policy to maximize expected return \(J(\theta)\).

At time step \(t\) in each episode:

1. Generate Action & Observe: Generate \(a_t\) following \(\pi(a|s_t, \theta_t)\), then observe reward \(r_{t+1}\) and next state \(s_{t+1}\).
2. Calculate Advantage (TD Error): Compute the TD error using the Critic's current value estimates: \[ \delta_t = r_{t+1} + \gamma v(s_{t+1}, w_t) - v(s_t, w_t) \]
3. Actor Update (Policy): Update the policy parameters to encourage actions with high advantage: \[ \theta_{t+1} = \theta_t + \alpha_{\theta} \delta_t \nabla_{\theta} \ln \pi(a_t | s_t, \theta_t) \]
4. Critic Update (Value): Update the value function parameters to minimize the TD error: \[ w_{t+1} = w_t + \alpha_w \delta_t \nabla_w v(s_t, w_t) \]

\[ \nabla_{\theta} J(\theta) = \mathbb{E}_{S \sim \rho, A \sim \beta} \left[ \frac{\pi(A|S, \theta)}{\beta(A|S)} \nabla_{\theta} \ln \pi(A|S, \theta) q_{\pi}(S, A) \right] \]

The update is similar to A2C but scaled by the importance sampling ratio \(\rho_t\):

\[ \theta_{t+1} = \theta_t + \alpha \underbrace{ \frac{\pi(a_t|s_t, \theta_t)}{\beta(a_t|s_t)} }_{\text{Importance Weight } \rho_t} \delta_t(s_t, a_t) \nabla_{\theta} \ln \pi(a_t \mid s_t, \theta_t) \]

Because of the log-derivative trick: \[ \theta_{t+1} = \theta_t + \alpha \frac{\delta_t(s_t, a_t)}{\beta(a_{t}|s_{t})} \nabla_{\theta} \pi(a_t \mid s_t, \theta_t) \]

Algorithm is similar only with difference of factor \(\beta(a_{t}|s_{t})\)

We write the probability of a trajectory with MDP for \(\tau = (s_0, a_0, \dots, s_T, a_T)\) :

• \(\tau\): Trajectory
• \(\pi_{\theta}\): Policy parameterized by \(\theta\)
• \(P(s_1)\): Probability of the initial state
• \(\pi_{\theta}(a_t | s_t)\): Probability of taking action \(a_t\) in state \(s_t\) (The Policy)
• \(P(s_{t+1} | s_t, a_t)\): Probability of transitioning to \(s_{t+1}\) given \(s_t\) and \(a_t\) (The Transition Function/Model)

Taking the gradient of the log-probability: \[\nabla_\theta \log P(\tau|\pi_{\theta}) = \nabla_\theta \log \rho_0(s_0) + \sum_{t=0}^{T} \left( \nabla_\theta \log P(s_{t+1}|s_t, a_t) + \nabla_\theta \log \pi_\theta(a_t|s_t) \right)\]

Since the initial state distribution \(\rho_0\) and the environment dynamics \(P(s_{t+1}|s_t, a_t)\) do not depend on the policy parameters \(\theta\), their gradients are zero: \[\nabla_\theta \log P(\tau|\theta) = \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t)\]

We aim to maximize the objective \(J(\pi_\theta) = \mathbb{E}_{\tau \sim \pi_\theta} [R(\tau)]\).

\[\nabla_\theta J(\pi_\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \left( \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) \right) R(\tau) \right]\]

From the above function, we introduce a baseline for \(R{\tau}\) for reducing the variance. A great option is the current value \(V(\tau)\). So we have the advantage function for objective function derivative: A = Q - V.

• Monte Carlo: using G to estaminate Q: A = G - V
• Temporal Difference: using \(Q=r+ \gamma V(s_{t+1}​)\), so \(A = \delta_t^{V} = Q-V = r_t + \gamma V(s_{t+1}) - V(s_t)\)
• GAE: MC cover all the steps, and TD only see one step, so GAE use l steps for generall case \(A_t^{GAE} = \sum_{l=0}^{\infty} (\gamma \lambda )^l \delta_{t+l}^V\)

use MSE, \((V-Q)^2\), where V is the output of value model and Q is the expected value of current step, we use Q = A + V to share the advantage calculation from above.

The value function (critic) is trained by minimizing the mean squared error between the predicted value and the estimated return:

\[ \mathcal{L}_{\text{value}} = \mathbb{E}_{t} \left[ \left( V_\psi(s_t) - ( A_t + V_\psi(s_t) ) \right)^2 \right] \]

Unlike traditional RL (e.g., Atari) where state representation is efficiently shared, LLM-based RL often requires decoupling the Actor and Critic due to Negative Transfer and Optimization Divergence.

import torch
import torch.nn.functional as F

# constants
kl_beta = 0.1
critic_weight = 0.5
ppo_eps = 0.2

# sample prompt completions and rewards
with torch.no_grad():
    completions = LLM.generate(prompts)  # (B*G, L)
    rewards = RM(completions)  # (B*G, 1)

# create a padding mask from lengths of completions in batch
completion_mask = <... mask out padding tokens ...>

# compute value function / critic output
values = CRITIC(completions)  # (B*G, L) - predicted reward per token!

# get policy logprobs for each action
llm_out = LLM(completions)
per_token_logps = F.log_softmax(llm_out, dim=-1)  # (B*G, L)

# get reference logprobs for each action
ref_out = REF(completions)
ref_per_token_logps = F.log_softmax(ref_out, dim=-1)  # (B*G, L)

# compute KL divergence between policy and reference policy
kl_div = per_token_logps - ref_per_token_logps

# directly subtract KL divergence from rewards
# NOTE: KL div is per token, so reward becomes per token and reward
# for all tokens (besides last token) is just kl divergence.
# Reward for last token is sum of outcome reward and KL div.
rewards -= kl_beta * kl_div # (B*G, L)

# compute the advantage - simple approach
advantage = rewards - values.detach()  # (B*G, L)

# compute the policy ratio
# NOTE: old_per_token_logps must be persisted during first policy
# update for this batch of data and re-used in each subsequent update
policy_ratio = torch.exp(
    per_token_logps - old_per_token_logps,
)  # (B*G, L)
clip_policy_ratio = torch.clamp(
    policy_ratio,
    min=1.0 - ppo_eps,
    max=1.0 + ppo_eps,
)

# compute the ppo loss
ppo_loss = torch.min(
    advantage * policy_ratio,
    advantage * clip_policy_ratio,
)  # (B*G, L)
ppo_loss = -ppo_loss

# combine ppo loss and critic mse loss
critic_loss = ((rewards - values) ** 2)  # (B*G, L)
loss = ppo_loss + critic_weight * critic_loss

# aggregate the loss across tokens (many options exist here)
loss = ((loss * completion_mask).sum(axis=-1) /
        completion_mask.sum(axis=-1)).mean()
     
# perform policy gradient update
optimizer.zero_grad()
loss.backward()
optimizer.step()

Initialize the policy model, value model, optimizer, and freeze a reference policy model

Sample a prompt from the dataset to start a rollout

Prompt + current policy model generate a completion (no gradient)

Completion + reward model produce a scalar reward per token or per sequence

Completion + current policy model produce current action log-probabilities

Completion + frozen reference policy model produce reference log-probabilities

KL divergence between current and reference log-probabilities is computed

The KL penalty is added to the reward to form the final shaped reward

Completion + value model produce value estimates for each timestep

Shaped rewards and value estimates are combined using GAE to compute advantages and returns

Stored rollout data are shuffled and split into minibatches for PPO training

Prompt and completion are passed again through the current policy to compute updated log-probabilities

The ratio between updated policy probabilities and stored old policy probabilities is computed

The clipped PPO objective uses this ratio and the advantages to compute the policy (actor) loss

The value model is trained using mean-squared error between predicted values and computed returns

Policy loss, value loss, and entropy bonus are summed to form the total PPO loss

Backpropagation updates both policy and value model parameters in one optimizer step

Steps 11–16 are repeated for multiple PPO epochs over the same rollout data

A new rollout is collected using the updated policy, and the process repeats

for prompt in dataset:
    # 1. Rollout
    generated_tokens = ActorModel.generate(prompt)
    full_text = prompt + generated_tokens

    # 2. Final Reward Score
    reward_score = RewardModel(full_text)

    deltas = []

    # 3. Token-Level Loop (Iterate over GENERATED tokens only)
    for t in generated_tokens:

        # KL Penalty
        prob_actor = ActorModel(prompt + tokens_up_to_t)
        prob_ref = ReferenceModel(prompt + tokens_up_to_t)
        KL_t = log(prob_actor) - log(prob_ref)

        # Step Reward
        if t == last_token:
            r_t = reward_score - (beta * KL_t)
        else:
            r_t = -(beta * KL_t)

        # TD Error (Delta)
        V_current = CriticModel(prompt + tokens_up_to_t)
        V_next = CriticModel(prompt + tokens_up_to_t + t) 
        delta_t = r_t + (gamma * V_next) - V_current

        deltas.append(delta_t)

    # 4. Calculate GAE (Advantage)
    # Usually computed in reverse order to properly apply (gamma * lambda) discounts
    A_t = compute_gae_backwards(deltas, gamma, lambda)

    # 5. PPO Loss Optimization
    # 'ratio' is the probability of the token under updated weights vs old weights
    ratio = P_actor_new(t) / P_actor_old(t) 

    # Clip the ratio, multiply by Advantage, and maximize (or minimize negative)
    L = -mean( min(ratio * A_t, clip(ratio, 1-epsilon, 1+epsilon) * A_t) )

    Loss.backward()
    Optimizer.step()




---------------------------------------------
for batch_prompts in dataset:

    # ==========================================
    # PHASE 1: ROLLOUT (No Gradients)
    # ==========================================
    with torch.no_grad():
        # 1. Generate responses
        responses = ActorModel.generate(batch_prompts)
        full_texts = batch_prompts + responses

        # 2. Get old log probabilities and values
        old_log_probs = ActorModel.get_log_probs(full_texts)
        ref_log_probs = ReferenceModel.get_log_probs(full_texts)
        values = CriticModel(full_texts)

        # 3. Get Reward Model scores for the finished sentences
        reward_scores = RewardModel(full_texts)

        # 4. Vectorized step math (Calculated all at once, no loop!)
        KL_penalties = old_log_probs - ref_log_probs
        step_rewards = compute_step_rewards(reward_scores, KL_penalties)

        # 5. Compute GAE and Returns for the whole batch
        advantages, returns = compute_gae_vectorized(step_rewards, values, gamma, lambda)

    # ==========================================
    # PHASE 2: PPO TRAINING (With Gradients)
    # ==========================================
    # PPO reuses the rollout data for multiple epochs
    for ppo_epoch in range(PPO_EPOCHS):
        for mini_batch in create_mini_batches(full_texts, old_log_probs, advantages, returns):

            # 1. Get NEW log probs and NEW values from updated models
            new_log_probs = ActorModel.get_log_probs(mini_batch.texts)
            new_values = CriticModel(mini_batch.texts)

            # 2. Calculate Ratio
            ratio = torch.exp(new_log_probs - mini_batch.old_log_probs)

            # 3. ACTOR LOSS (Policy Loss with Clipping)
            surr1 = ratio * mini_batch.advantages
            surr2 = torch.clamp(ratio, 1.0 - epsilon, 1.0 + epsilon) * mini_batch.advantages
            actor_loss = -torch.min(surr1, surr2).mean()

            # 4. CRITIC LOSS (Value Loss: MSE between predicted values and actual returns)
            critic_loss = MSE(new_values, mini_batch.returns)

            # 5. Total Loss & Backprop
            total_loss = actor_loss + (value_coefficient * critic_loss)

            total_loss.backward()
            Optimizer.step()

We reject GRPO (which requires generating groups for a baseline). Instead, we use Time as our baseline.

We use a single Transformer. The MTP Modules (which normally predict future tokens) are slightly modified to also predict the Value (Expected Future Reward) of those tokens.

• No Separate Critic Network: The MTP heads are the Critic.
• No Group Sampling: We do not compare against a group average. We compare against our own prediction from the next step (Bootstrapping).
• One Network: Parameters \(\theta\) are shared.

Standard DeepSeek MTP predicts tokens \(t_{n+1}, t_{n+2} \dots\). We modify the output projection of the MTP modules to output two things:

1. Token Logits: \(P(t_{n+k})\) (What happens next?)
2. Scalar Value: \(V_{n+k}\) (How good is it?)

Equation: \[ [ \text{Logits}_{k}, v_{k} ] = \text{MTP Head}_k(h_n) \] Where \(v_k\) represents the expected return starting from step \(n+k\).

1. Single Network: The Policy and Value are fused. The "Value" is just a tiny scalar output on the existing MTP heads.
2. No GRPO: We don't need multiple samples to find a baseline. We use the Bellman Consistency (\(V_t \approx r + V_{t+1}\)) as the training signal.
3. MTP Integration: The MTP heads are essential. They provide the "Lookahead" capability that stabilizes the single-network value estimation (reducing the noise of a single step).
4. Internal Value Saved: The training relies on carrying the scalar \(v\) from step \(t+1\) backward to step \(t\).

A unified architecture where MTP Heads serve a dual purpose:

1. Syntax/Logic: Predicting future tokens (Language Modeling).
2. Planning: Predicting future value (Implicit Critic).

• Extreme Efficiency: Eliminates the memory cost of a separate Critic (PPO) and the compute cost of Group Sampling (GRPO).
• Temporal Credit Assignment: Unlike GRPO (which gives the same reward to the whole sentence), this method assigns specific values to specific tokens via bootstrapping.
• Dopamine Signals: The MTP value prediction (\(v_t\)) acts like a localized dopamine signal, telling the model exactly when it made a good move, not just at the end.

Since we use the same network to generate the target \(V_{t+1}\) and the prediction \(V_t\), the training can oscillate or diverge.

We cannot easily afford a second full network. However, we can keep a lightweight copy of just the MTP heads (a few MBs).

• Active Heads (\(\theta\)): Learn rapidly.
• Target Heads (\(\theta'\)): Update slowly (\(\theta' \leftarrow \tau\theta + (1-\tau)\theta'\)).
• Stabilized Rule: \[ Y_t = r_{t+1} + \gamma \text{MTP}_{\text{target}}(s_{t+1}) \]

This algorithm is feasible but requires careful tuning of the Auxiliary Loss Balance.

\[ L = L_{\text{token}} + \lambda_1 L_{\text{MTP\_tokens}} + \lambda_2 L_{\text{TD\_Value}} \]

• If \(\lambda_2\) is too high, the "Value" objective will overwrite the "Language" objective (Catastrophic Forgetting).
• If \(\lambda_2\) is too low, the agent won't plan.
• Recommendation: Start with \(\lambda_2 = 0.1\) and clamp the value gradients.