Unified RL

All matrix math here will follow pytorch conventions. I.E. applying a matrix $M$ to a vector $x$ will look like $xM$ rather than $Mx$. Further applying a matrix $W$ will hence look like $xMW$ rather than $WMx$.

1. Intro

1.1. Definitions

Markov State ($S$)

An information state ($S_t$) is considered Markov iff it contains all useful information from the history:

\[P\left(S_{t + 1} \mid S_t \right) = P \left(S_{t + 1} \mid S_1, \cdots, S_{t}\right)\]

Markov states are desirable because agents are generally expected to be long-lived and take very frequent actions. The analogy here is with LLMs. LLM states are not exactly markov, you have to keep the entire KV-cache, and that’s your “history”. The could in theory be markov, but the residual stream would just have to be much larger to contain the expressivity of the KV cache.

Markov Decision Process (MDP)

A process where the agent observes the full environment state.

Partially Observable MDP (POMDP)

A process where the agent indirectly observes the environment state (i.e. there are hidden variables controlling environment behavior). E.g. stock trader only sees stock prices, but doesn’t know driving forces.

Policy ($\pi$)

A function from state $\rightarrow$ actions. There are:

• Determistic Policies
• Stochastic Policies

Value Function ($V$)

“Expected total future reward”.

\[v_\pi(s) = \mathbb{E}_\pi \left[ R_t + \gamma R_{t + 1} + \gamma^2 R_{t + 2} + \cdots \mid S_t = s \right]\]

Model

Something that predicts what the environment will do next. Typically this includes:

A Transition Model / table ($\mathcal{P}$)

\[\mathcal{P}_{s \rightarrow s'}^a = P \left[ S' = s' \mid S = s, A = s \right]\]

A Reward Model (this below is rather straightforward)

\[\mathcal{R}_s^a = \mathbb{E} \left[ R \mid S = s, A = a \right]\]

A model is not required. Many techniques are model-free.

1.2. Types of RL Agents

The “type” basically describes which of the key components {policy, value_function, model} the agent uses.

• Value Based: No explicit policy (e.g. “greedy”); uses Value Function
• Policy Based: No explicit value function, just a policy table (e.g. “if at this square in a maze, go right”).
• Actor Critic: Has both value function and policy.
• Model Free: orthogonally from Value Function / Policy, we choose to ignore trying to build out a model of how the environment works.

1.3 Types of RL Problems

RL:

• Multi-rollout (agent gets to learn from unlimited rollouts) in same environment
• Multi-rollout in changing environment (structure of environment is same)
• Single-rollout

2. Introducing MDPs

Most RL problems can be converted to MDPs. Let’s take a look at an MDP where we happen to have a fully observable environment.

As it turns out, having a fully observable environment is not make-or-break, as partially observable RL problems can be converted to POMDPs, which can be converted to MDPs.

2.1. Model / Transition table ($\mathcal{P}$)

Since we have a fully observable environment, then we can build a transition matrix that models the probabilities we believe each state ($s$) will lead to a successive state ($s’$).

\[\begin{align*} \because S' &= S \mathcal{P} \\ \therefore \mathcal{P} &= \begin{bmatrix} \mathcal{P}_{s_1 \rightarrow s_1} & \cdots & \mathcal{P}_{s_1 \rightarrow s_n} \\ \vdots & \ddots & \vdots \\ \mathcal{P}_{s_n \rightarrow s_1} & \cdots & \mathcal{P}_{s_n \rightarrow s_n} \end{bmatrix} \end{align*}\]

Each row of $\mathcal{P}$ sums to 1.

In an MDP where there are various actions, each action $a \in \mathcal{A}$ will yield one of these 2D matrices, which we call $\mathcal{P}^a$.

2.2. Markov Reward Process

A Markov Reward Process is simply an “Markov Chain / rollout with rewards”, i.e. a rollout of an MDP where there are no chosen actions. Everything is automatic. This can be written / parameterized as $\langle \mathcal{S}, \mathcal{P}, \mathcal{R}, \gamma \rangle$, where:

• $\mathcal{S}$ is the set of states
• $\mathcal{P}$ is the transition table
• \(\mathcal{R}_s\) is a reward function (e.g. \(R_t = \mathbb{E} [R_{t + 1} \mid S_t = s]\)). Note:
  • $R_{t + 1}$ is the reward for taking actions at time $t$
  • Conversely, $R_t$ is the reward for simply having arrived at the current state at time $t$.
• $\gamma \in [0, 1]$ is the discount factor. We do this because:
  • We have uncertainty about our forecasting ability
  • To keep the math bounded. You can’t have infinite returns.

Return

The “return” $G_t$ is the total discounted reward from time-step $t$, computed as you are ABOUT TO make the next step

\[\begin{align*} G_t &= R_{t + 1} + \gamma R_{t + 2} + \gamma^2 R_{t + 3} + \cdots \\ &= \sum_{k = 0}^\infty \gamma^k R_{k + t + 1} \end{align*}\]

Value Function

\[v(s) = \mathbb{E}[G_t \mid S_t = s]\]

This is a statement of fact. Its a higher-level abstraction that does not concern itself with how $G_t$ was gotten (i.e. what policy). When we concern ourselves with the “decision”-making part of “MDP”, we will decide on some policy to follow, but the value function is still conceptually this.

Because $v(s)$ is an expectation, this implies that you have to sample rollouts so that you have multiple rollouts and returns to take the expectation of.

Bellman Equation for $v$ in MRP

The Bellman Equation is just a recursive formulation of $v$:

\[\begin{align*} v(s) &= \mathbb{E}[G_t \mid S_t = s] \\ &= \mathbb{E}[R_{t + 1} + \gamma (R_{t + 2} + \gamma R_{t + 3} + \cdots) \mid S_t = s] \\ &= \mathbb{E}[R_{t + 1} + v(S_{t + 1}) \mid S_t = s] \end{align*}\]

Factoring in $\mathcal{P}$:

\[\begin{align*} v(s) = \mathcal{R}_s + \gamma \mathcal{P}_{[s,:]} \cdot \begin{bmatrix} v(s_1) \\ \vdots \\ v(s_n) \end{bmatrix} \end{align*}\]

Visually, we have:

Bellman in Matrix Form

Let:

\[\begin{align*} v = \begin{bmatrix} v(s_1) \\ \vdots \\ v(s_n) \end{bmatrix}, \text{ and } \mathcal{R} = \begin{bmatrix} \mathcal{R}_{s_1} \\ \vdots \\ \mathcal{R}_{s_n} \end{bmatrix} \end{align*}\]

Then

\[\begin{align*} v & = \mathcal{R} + \gamma \mathcal{P} v \end{align*}\]

This means there exists a closed form solution for $v$, which is simply:

\[\begin{align*} v & = (I - \gamma \mathcal{P})^{-1} \mathcal{R} \end{align*}\]

The only thing is that the inverse has $O(n^3)$ computational complexity, so it’s not typically practical for large MRPs. So in practice, there exist a few different ways to solve this that are all variants of Dynamic Programming:

• Simple Dynamic Programming
• Monte-Carlo Evaluation
• Temporal-Difference Learning

2.3. Markov Decision Process

An MDP is the same as an MRP, just that we now have agency (i.e. a set of actions to choose from). It is represented as $\langle \mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma \rangle$, where:

• $\mathcal{A}$ is the set of actions.
• \(\mathcal{R}_{s}^a\) is the reward for being at state \(s\) and deciding to take action \(a\). Sometimes interchangable with \(R_{t + 1}\).

Policy

A policy $\pi$ is just a distribution over possible actions to take, given a state:

\[\begin{align*} \pi(a \mid s) & = P\left[ A_t = a \mid S_t = s \right] \end{align*}\]

Value Function

For an MRP, we had:

\[v(s) = \mathbb{E}[G_t \mid S_t = s]\]

This does not have any sense of action-taking baked into it. The process just unfolds by itself without any choice of actions. Now, we have a policy that tells us how to take action. $v$ becomes a distribution (because $\pi$ is a distribution over actions), and we can then take the expectation of it.

Bellman Equation for $v$ in MDP

For an MRP, $v$ has the Bellman Equation:

\[\begin{align*} v & = \mathcal{R} + \gamma \mathcal{P} v \end{align*}\]

If we were to add in the action-taking steps, we arrive at this decision tree:

We are at state $\boldsymbol{s}$. From there, we sample a next action \(A_{t + 1}\), which transitions into an unknown new state \(S_{t + 1}\), and so on. So, the expectation becomes:

\[\begin{align*} v_\pi(s) &= \sum_{a \in \mathcal{A}} \pi(a \mid s) \left( \mathcal{R}_s^a + \gamma \mathcal{P}_{[s,:]}^a \cdot v_\pi \right) \end{align*}\]

Action-Value Function ($q$)

So far, we’ve talked about only the state value function $v$. Now, we can introduce the action-value function $q(s, a)$:

\[\begin{align*} q_\pi (s, a) &= \mathbb{E}_\pi [G_t \mid S_t = s, A_t = a] \end{align*}\]

Bellman Equation for $q$

\[\begin{align*} q_\pi(s, a) &= \mathbb{E}_\pi[G_t \mid S_t = s, A_t = a] \\ &= \mathbb{E}_\pi[R_{t + 1} + \gamma (R_{t + 2} + \gamma R_{t + 3} + \cdots) \mid S_t = s, A_t = a] \\ &= \mathbb{E}_\pi[R_{t + 1} + \gamma q_\pi(S_{t + 1}, A_{t + 1}) \mid S_t = s, A_t = a] \end{align*}\]

But we’re not actually done yet, because we’re being imprecise about what \(R_{t + 1}\), \(S_{t + 1}\), and \(A_{t + 1}\) are.

The bolded s and a_1 nodes are the specified operands in q(s, a). Everything else after is unknown.

As you can see from the above decision tree diagram, we only know what $\boldsymbol{s}$ and $a$ are in \(q_\pi(s, a)\), hence \(R_{t + 1} = R_s^{a = a_1}\). The continuation of the path from there is a random variable, hence \(S_{t + 1}\) and \(A_{t + 1}\) are random variables. To fill in this piece of information to make the Bellman Equation more precise, we have:

\[\begin{align*} q_\pi(s, a) &= R_s^a + \gamma \sum_{s' \in \mathcal{S}} \mathcal{P}_{s \rightarrow s'}^a \underbrace{\sum_{a' \in \mathcal{A}} \pi(a' \mid s') q_\pi(s', a')}_{\textstyle = v_\pi(s')} \end{align*}\]

Connection between $q$ and $v$

By now, we can see that $q$ and $v$ are just two sides of the same coin - one allows you to specify the next $a$, one does not.

3. Solving the Planning Problem

3.1. Optimizing $v$, $q$, and $\pi$

In sequential decision making problems, we want to maximize $v$ / $q$. We will simply optimize $q$ because $q$ is a more useful formulation of value.

\[\begin{align*} q_\star (s, a) &= \max_{\pi} q_\pi(s, a) \end{align*}\]

Optimal Policy $\pi_\star$

A policy $\pi$ is better than another policy $\pi’$ if:

\[\begin{align*} v_\pi (s) \geq v_{\pi'} (s), \forall s \end{align*}\]

There is a very important Theorem about MDPs: For any MDP, there exists an optimal policy \(\pi_\star\) that is better than or equal to all other policies \(\pi\). This is amazing because it means we don’t have to reason about where in an MDP a policy might be better or worse than another.

Optimal Action-Value Function Q*

The optimal (deterministic) policy can be found by maximizing $q_\star$, and the optimal is hence a greedy one:

\[\begin{align*} \pi_\star (a \mid s) & = \begin{cases} 1 & \text{if } a = \operatorname*{argmax}\limits_{a} q_\star(s, a) \\ 0 & \text{otherwise} \end{cases} \end{align*}\]

Bellman Optimality Equation for $v_\star$

Previously, we arrived at the Bellman Equation for $v$ in a general MDP as:

\[\begin{align*} v_\pi(s) &= \sum_{a \in \mathcal{A}} \pi(a \mid s) \left( \mathcal{R}_s^a + \gamma \mathcal{P}_{[s,:]}^a \cdot v_\pi \right) \end{align*}\]

Notice that in a deterministic optimal policy, our actions are greedy. There’s no variability in the actions we choose - they are always the value-maximizing action, so we have:

\[\begin{align*} v_\star(s) &= \max_a \left( \mathcal{R}_s^a + \gamma \mathcal{P}_{[s,:]}^a \cdot v_\star \right) \end{align*}\]

Bellman Optimality Equation for $q_\star$

Similarly, we previously arrived at this Bellman Equation for $q$ in a general MDP:

\[\begin{align*} q_\pi(s, a) &= R_s^a + \gamma \sum_{s' \in \mathcal{S}} \mathcal{P}_{s \rightarrow s'}^a \underbrace{\sum_{a' \in \mathcal{A}} \pi(a' \mid s') q_\pi(s', a')}_{\textstyle = v_\pi(s')} \end{align*}\]

Factoring in the action greediness, we have:

\[\begin{align*} q_\star(s, a) &= R_s^a + \gamma \sum_{s' \in \mathcal{S}} \mathcal{P}_{s \rightarrow s'}^a \underbrace{\max_a q_\star(s', a')}_{\textstyle = v_\star(s')} \end{align*}\]

3.2. Solving Bellman Optimality Equations

In an MRP, we saw that we could solve the Bellman Equation by simple linear algebra, because in an MRP, we simply took expectations over actions and transitions - everything was a linear equation. In an MDP where you have a greedy policy, the greediness introduces a non-linearity, so there no longer exists a closed form solution (not that the closed form solution would have been practical anyway).

So, we have to use iterative methods, including:

• Value Iteration
• Policy Iteration
• Q-Learning
• SARSA

3.3. Policy Iteration

Policy Evaluation

Given an MDP \(\langle \mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma \rangle\), and a policy $\pi$, we can evaluate the policy (calculate $v_\pi$). Since this works for any policy, even non-greedy ones, we’ll use the Bellman expectation equation instead of the optimality equation:

\[\begin{align*} v_\pi(s) &= \sum_{a \in \mathcal{A}} \pi(a \mid s) \left( \mathcal{R}_s^a + \gamma \mathcal{P}_{[s,:]}^a \cdot v_\pi \right) \end{align*}\]

The formulation here already directly tells us how to update our $v$ values:

\[\begin{align*} v_{k + 1}(s) &= \sum_{a \in \mathcal{A}} \pi(a \mid s) \left( \mathcal{R}_s^a + \gamma \mathcal{P}_{[s,:]}^a \cdot v_{k} \right) \end{align*}\]

What we do is just to traverse all the $v$ values and update all of them per timestep ($k$).

Example:

Here’s the setup: a 🦆 (our agent) lives on a 4×4 grid and wants to reach a checkered flag 🏁. The two flag cells are terminal — the episode ends once the duck steps on either. On every other cell the duck can move up / down / left / right (moves off the edge leave it in place), and each step costs a reward of −1. Under a uniform-random policy, how good is each cell?

🏁

🦆

🏁

🏁 = terminal goal states; 🦆 = the agent.
Reward: −1.0 per step
γ: 1.0
Transition Prob: 1.0
π: uniform-random

Policy Evaluation in action:

k = 0

Value Iteration: v_k slowly converges to v_π

Policy Improvement

So now, having evaluated the policy to arrive at converged values of each state, we can use those values to infer a policy that is optimal (w.r.t to these values). In the 4x4 grid example, the optimal policy should just tell us to proceed in the direction of most reward (least red neighbor).

Greedy policy w.r.t. the converged v_π

Policy Iteration = Evaluation + Improvement

Ping-ponging between Evaluation and Improvement gives us Policy Iteration. Crucially, you don’t have to run the Evaluation step to convergence each time. You can see that even with 5 time-steps of Evaluation, we can already sufficiently update $v$ such that we can meaningfully improve on the policy (which happens to already be optimal):

Greedy policy w.r.t. v₅

The reason we can do partial evaluations and still update the policy is because of the powerful theorem that this method of evaluation and improvement will always converge to the optimal deterministic policy. The proof for this is simply the argument along the lines:

\[q_\pi (s, s(\pi')) = \max_{a \in \mathcal{A}} q_\pi (s, a) \geq \left(q_\pi (s, \pi(s)) = v_\pi (s) \right)\]

English: if we acted greedily for this 1 step from $s$ but acted according to $\pi$ from there-on, we will do at least as well as before. If the LHS and RHS are the same, then at least for the state, $\pi$ is optimal, and we have satisfied the Bellman Optimality Equation. The intuition for why this improvement will always lead to the optimal policy is non-obvious to me, but I find it helpful to remember that all states are updated. This isn’t the case of true Reinforcement Learning problems where we only have information on states that we visit (sample) and can hence be stuck in some local optimum.

Implications for Implementation

Stopping Conditions for Evaluation Step - we can:

• Updates for values are within some $\epsilon$
• Simply stop after $K$ iterations
• Extreme: stop after $K = 1$. This is Value Iteration

3.4. Value Iteration

The crux of Value Iteration is the same as in Policy Iteration: repeatedly using the Bellman Optimality (not Expectation this time) Equation to update the value scores for every state. The key differences are:

• You don’t have an explicit policy to follow - everything subsequent state can be chosen greedily during Evaluation. In the above Policy Iteration example, we followed the uniform-random policy and applied the Bellman Expectation Equation during Evaluation.
• For each round of value updates, you do only 1 sweep over all the states ($K = 1$)

In practice, you almost always us Value Iteration, unless there is an explicit policy to follow. Value Iteration can be thought of as the special case of Policy Iteration where your policy is simply greedy.

3.5. Implementation Tricks

Because the convergence conditions for Policy / Value Iteration are (1) greedy update and (2) you continually update all states (no states left behind), you can speed up the training process without affecting convergence possibility by:

• Doing In-Place updates of values
• Parallelizing updates over all states
• Asynchronously sampling and updating states
• Sampling rollouts and update only information seen during the rollout

These are all techniques used in large RL problems, where state spaces can be huge (and hence updating all of them always can sometimes be a waste of time), and some techniques also introduce stochasticity, which can be useful at times.

4. Model-Free Planning

In the previous methods where we could use variants of Dynamic Programming (Policy Iteration, Value Iteration) to solve MDPs, we could do this because:

• Transition Dynamics were known (we had the model $\mathcal{P}$). In our examples, the transitions were all deterministic, but the methods would have worked too for stochastic transitions.
• The Rewards $\mathcal{R}$ were known as well.
• The problem was not large (millions of states or smaller)

In the real world, either they may be trivial (deterministic), or more crucially, we may not have $\mathcal{P}$ or $\mathcal{R}$. However, what matters is that we are able to estimate the expectation of these. Sampling trajectories allows us to still do so. Goodbye models.

4.1. Sampling Full Episodes: Monte-Carlo RL

Suppose you have a simple trajectory that looks like this:

🦆

There are two variants of Monto-Carlo Policy Evaluation (computing $v_\pi(s)$).

1. “First-Visit”: We only use information collected at each state the first time that state is visited. That means the information collected at repeated visits (when the cell above is annotated by a second inset ring) is discarded.
2. “Every-Visit”: We use all information at every visit.

It is not clear which is better, but intuitively, I would think “Every-Visit” makes more sense, because of the Markov Property assumption. Moving forward, I’ll use this paradigm.

4.1.1. Collecting Episodes, Computing Empirical Means

Our goal here is to estimate \(v_\pi(s)\), which is an expectation of returns:

\[\begin{align*} v_\pi(s) &= \mathbb{E}_\pi[G_t \mid S_t = s] \end{align*}\]

So let’s walk through the above example to demonstrate the procedure for doing so. In the above example, we have this episode:

If this were the start of the learning process, for the states that are visited once, we simply set $v(s_t) = G_t$:

However, if it weren’t the start of the learning process, there may already be several prior episodes where we saw the same state featured:

Previous Episodes:

So this means that you have to keep a running count of how many times you’ve updated $v(s)$ for each $s$, so that you can scale updates properly such that $v(s)$ always reflects the mean. Specifically, this means that everytime you see an episode suffix beginning with state $s$, you do:

1. $n(s) \leftarrow n(s) + 1$
2. $v(s) \leftarrow v(s) + \frac{1}{n(s)}(G - v(s))$

In non-stationary problems, it can sometimes also be useful to just prioritize more recent values and let old values decay, so we just do:

• $v(s) \leftarrow v(s) + \alpha(G - v(s))$

4.2. Sampling Actions: Temporal-Difference RL (TD Learning)

TD-Learning is conceptually the same as MC Learning: it’s model-free, so we use sampling to try and estimate the mean value of states. The one key difference is that MC uses full trajectories / episodes, while TD-Learning uses current value estimates of successor states, just like in the Bellman Equation for MDPs. This means that TD-Learning can also be used in non-terminating environments, whereas MC can only be applied in episodic environments.

4.2.1. TD(0)

\[\begin{align*} V(S_t) \leftarrow V(S_t) + \alpha (\underbrace{\underbrace{R_{t + 1} + \gamma V(S_{t+1})}_{\textstyle \text{TD target}} - V(S_t)}_{\textstyle \delta_t = \text{TD error}}) \end{align*}\]

Pictorially:

4.3. MC vs TD

• MC has better convergence properties (it will converge to the true values even with function approximators for $V$)
• MC does not assume the Markov property, whereas TD does (TD implicitly builds an MDP)
• TD is usually much more efficient

4.4. MC vs TD vs DP

So far we’ve seen Dynamic Programming (Policy / Value Iteration), MC Learning, and TD Learning. They differ along 2 axes:

• Sampling vs no sampling (“full backups”)
• Shallow vs Deep backups

4.5. TD($\lambda$)

TD($\lambda$) is a way of interpolating between TD(0) and MC. $\lambda$ specifies how many real steps to take, before imputing the rest with current estimates of $V$. You can imagine that TD($\infty$) is equivalent to MC.

$\lambda$-return

There is a concept known as $\lambda$-return, where we try to take the geometric mean between the returns computed via TD(0), TD(1), …, TD($\infty$). This does not seem immediately useful, so I’m skipping this.

5. How To Sample

In the previous section, we arrived at a clear idea of how we may use sampled episodes / trajectories to update our $v$ values. In this section, we’ll explore how those episodes are sampled. After all, in real RL problems, the experience we get is only what the agent has actually done.

5.1. The Problem with $V$.

Previously, we already had the episode / partial episode samples. We already knew which state \(S_{t + 1}\) followed the current state \(S_t\). When we’re generating these trajectories, we need to know how to act. Most of the time, our action will be greedy, which means we need to take:

\[\begin{align*} a = \operatorname*{argmax}\limits_{a \in \mathcal{A}} \left(R_{s + 1} + \gamma \mathbb{E} [V_{s + 1}]\right) \end{align*}\]

The problem with this is that \(\mathbb{E} [V_{s + 1}]\) requires a model:

\[\begin{align*} \mathbb{E} [V_{s + 1}] = \mathcal{P}_{[s:]}^a \cdot v \end{align*}\]

And we don’t have a model! So to solve this, we use $Q$ instead. This way, we can offload the expectation part into $Q(s, a)$.

5.2. Explore vs Exploit

The next idea is leaving room for exploration. If you only exploit the best action you’ve seen so far, you’re susceptible to never knowing any better option. Therefore, we leave some room for exploration in the policy:

\[\begin{align*} \pi(a \mid s) &= \begin{cases} \operatorname*{argmax}\limits_{a \in \mathcal{A}} Q(s, a) & \text{ with prob } 1 - \epsilon \\ \text{random } & \text{ with prob } \epsilon \end{cases} \end{align*}\]

Epsilon-greedy Q values and policy still do guarantee improvements (since Epsilon-greedy is just a mixture of greedy and uniform-random; the uniform-random bit improves iff $q$ values improves, which it does, since the other part is greedy)

Implementation Notes

In practice, we slowly decay $\epsilon \rightarrow 0$, so something like $\epsilon = 1/k$ where $k$ is the number of steps, or $\epsilon = \frac{1 + L}{k + L}$ where $L$ is some large integer.

5.3. Online TD-Learning: SARSA

Previously, we saw that TD-Learning has this update rule:

\[\begin{align*} V(S_t) \leftarrow V(S_t) + \alpha (\underbrace{\underbrace{R_{t + 1} + \gamma V(S_{t+1})}_{\textstyle \text{TD target}} - V(S_t)}_{\textstyle \delta_t = \text{TD error}}) \end{align*}\]

The above requires us to know \(S_t, R_{t + 1}, S_{t + 1}\). In replacing $V$ with $Q$, we now have:

\[\begin{align*} Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha (\underbrace{\underbrace{R_{t + 1} + \gamma Q(S_{t+1}, A_{t + 1})}_{\textstyle \text{TD target}} - Q(S_t, A_{t})}_{\textstyle \delta_t = \text{TD error}}) \end{align*}\]

Observe that this requires us to know \(S_t, A_t, R_{t + 1}, S_{t + 1}, A_{t + 1}\), hence the name SARSA. And similar to TD-$\lambda$, we can also interpolate between TD and MC in the $Q$ context by deciding how many actual steps to sample; this is known as “SARSA($\lambda$)”.

5.4. Off-Policy Learning

What if you wanted to evaluate your $V$ or $Q$ values according to $\pi$, but the episodes you have were generated while following some other behavior policy $\mu$? This is mostly a question of how to use already-existing data (from other robots, humans, etc.).

5.4.1. Importance Sampling

The problem with this is that obviously the distributions $\pi$ and $\mu$ are not the same, which means that your expectations are also not the same. Here, we introduce the idea of Importance Sampling, which applies scaling factors to the terms in the TD Update Equation so morph \(\mathbb{E}_{\sim \mu}\) into \(\mathbb{E}_{\sim \pi}\):

\[\begin{align*} V (S_t) \leftarrow V (S_{t + 1}) + \alpha \left( \frac{\pi (A_t \mid S_t)}{\mu (A_t \mid S_t)} \left( R_{t + 1} + \gamma V(S_{t + 1}) \right) - V(S_t) \right) \end{align*}\]

But obviously you may not know what $\mu(A_t \mid S_t)$ is. If you just found a dataset of trajectories lying around, you definitely won’t know what $\mu$ was.

5.4.2. Q-Learning

So we don’t use Importance Sampling, we use Q-Learning (a variant of SARSA)!

Unfortunately, Q-Learning doesn’t address the idea of how we can use found data generated from an unknown policy. I suspect you can use classic TD-learning to learn your $V$ values and back out a policy from there that you can then use to warm-start Q-learning on your own generated data, but this is out of scope.

What Q-Learning does is:

1. We sample an action \(A_{t + 1} \sim \mu(A_t \mid S_t)\) to actually take while exploring
2. We also sample an action \(A' \sim \pi(A_t \mid S_t)\) to use to do our Q updates.

🤖 follows the behaviour policy μ (sampling A_t+1); 🦆 marks the target policy π (sampling A′). The black path is the trajectory actually taken.

This allows us to explore off-$\pi$ while updating our $Q$ values according to $\pi$.

\[\begin{align*} Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha \left(R_{t + 1} + \gamma Q(S_{t + 1}, A') - Q(S_t, A_t) \right) \end{align*}\]

6. Function Approximators

Sometimes, RL problems can be huge, or states can be continuous. So what do we do? Use function approximators! This section will only cover deep networks since they are universal function approximators.

6.1. Formulation

Suppose our state $S$ is a vector (hence continuous), then the formulation is simply that $V$ is a deep network that outputs the estimate of \(V_\pi(S)\).

Our TD updates used to be this:

\[\begin{align*} V(S_t) \leftarrow V(S_t) + \alpha (\underbrace{\underbrace{R_{t + 1} + \gamma V(S_{t+1})}_{\textstyle \text{TD target}} - V(S_t)}_{\textstyle \delta_t = \text{TD error}}) \end{align*}\]

Everything is conceptually the same as before (we update towards the TD target), except that now:

• We use $\hat{V}_w$ instead of $V$ to denote that the value is just an estimate parameterized by $w$ (the weights of the deep network)
• We update $w$, which in turn updates the value generated by $V$

In ML parlance, we have:

Error

Something like MSE:

\[\begin{align*} \frac{1}{2} \left( R_{t + 1} + \gamma \hat{V}_w(S_{t+1}) - \hat{V}_w(S_t) ) \right)^2 \end{align*}\]

Gradient

\[\begin{align*} \left( R_{t + 1} + \gamma \hat{V}_w(S_{t+1}) - \hat{V}_w(S_t) ) \right) \nabla(- \hat{V}_w(s_t)) \end{align*}\]

Update

Negative Gradient $\times$ learning rate:

\[\begin{align*} \Delta w = \alpha \left( R_{t + 1} + \gamma \hat{V}_w(S_{t+1}) - \hat{V}_w(S_t) ) \right) \nabla(\hat{V}_w(s_t)) \end{align*}\]

A similar thing can be done for $Q$, but for $Q$, you have several options for the shape of $\hat{Q}$:

1. $Q: \mathbb{R}^\text{state dim} \rightarrow \mathbb{R}^\text{action dim}$ (useful for discrete action space)
2. $Q: \mathbb{R}^\text{state dim + action dim} \rightarrow \mathbb{R}^1$

6.2. Convergence & Efficiency

As you can see, not only are our $\hat{V}_w$ values changing, the way we compute those $\hat{V}_w$ values (i.e. $w$) is also changing (moving target). This can cause convergence troubles. Moreover, doing gradient descent after every sampled SARSA tuple is inefficient. Here’s where we solve both problems with Deep-Q Networks.

6.2.1. Deep-Q Networks (DQN)

In the below code, you’ll see:

• We have q_net_target, which is updated much less frequently than q_net. This solves the problem where $\hat{Q}_w$ is always moving (moving target problem).
• q_net_target is updated softly (interpolated between q_net and q_net_target).

Implementation

def soft_update_q_target(q_net, q_net_target, tau):
    """
    Desc:                   Architecture-agnostic helper function to copy weights.

    @param q_net:           (type[nn.Module]) The source net
    @param q_net_target:    (type[nn.Module]) The dest net
    """
    for param, target_param in zip(q_net.parameters(), q_net_target.parameters()):
        target_param.data.copy_(tau * param.data +
                                (1 - tau) * target_param.data)

q_net = QNet(...parameters)
q_net_target = copy.deepcopy(q_net)

for _ in range(num_train_episodes):
    episode_reward = 0.0
    s, _ = env.reset()
    for _ in range(max_steps_per_episode):
        # Step 1) we sample a transition to add to replay buffer.
        a = select_action_eps_greedy(env, q_net, s, eps)
        s_next, r = env.step(a)

        episode_reward += r
        replay_buffer.append((s, a, r, s_next)) # we will compute a_next on the fly

        # Step 2) sample a mini batch and fit the Q network
        if len(replay_buffer) < batch_size:
            continue

        batch_s, batch_a, batch_s_next, batch_r, batch_terminal = \
            sample_experience_batch(replay_buffer, batch_size)

        # Compute Q targets (in this case q_net goes from state_dim to action_dim)
        batch_s_next_q_vals = q_net_target(batch_s_next)
        batch_s_next_q_greedy = torch.max(batch_s_next_q_vals, axis=-1).values
        batch_targets = torch.where(
            batch_terminal,
            batch_r,
            batch_r + gamma * batch_s_next_q_greedy
        )

        # Update
        q_net.train(True)
        batch_s_q_vals = q_net(batch_s)[torch.arange(batch_size), batch_a]
        loss = F.mse_loss(batch_s_q_vals, batch_targets.detach())
        q_net_opt.zero_grad()
        loss.backward()
        q_net_opt.step()

        # Update q_net_target and evaluate
        if total_steps % q_target_update_freq == 0:
            soft_update_q_target(q_net, q_net_target, tau)

7. Policy Gradient

So far, everything we’ve seen revolves around Value-based methods, but Value-based methods may lose out to Policy-based methods because:

• Value-based methods require us to compute $V$ or $Q$ for our state, which may be much more tedious than a policy. For example, in pong, the value of a given state may be hard to predict, but the policy is simple: if ball is above your paddle, go up, and vice versa.
• Computing \(a = \operatorname*{argmax}\limits_{a} q_\star(s, a)\) in particular may also be hard because $\mathcal{A}$ may be huge.
• Policy-based methods allow you to specify explicit, stochastic policies that are more general than $\epsilon$-greedy.

Sections 7.1. to 7.5. basically speed through what metrics we might use for a policy, what naive methods there are, what better methods there are, and how we might implement those.

7.1. Policy Objective Functions

Given some policy $\pi_\theta$ parameterized by weights $\theta$, we want to optimize the policy. We first have to define how to measure the quality of $\pi_\theta$. For that, we have several options, including:

• Using the average start state value
• Using the average state value
• Using the average reward per time-step

In the above, “average” is w.r.t to the stationary distribution of the Markov chain for $\pi_\theta$.

7.2. Policy Gradients by Finite Distances

One way of estimating a policy gradient w.r.t to $\pi_\theta$ is by perturbing $\theta$ by small amounts. Below shows how you would perturb $\theta$ in the $k$-th dimension:

\[\begin{align*} \frac{\partial J(\theta)}{\partial \theta_k} \approx \frac{J(\theta + \epsilon \textbf{e}_k) - J(\theta)}{\epsilon} \end{align*}\]

7.3. Analytical Policy Gradients

There’s also an analytical way of arriving at a policy gradient, and that’s by treating the above metrics ($v$, $r$, $Q$, etc.) as constants w.r.t to $\theta$ (which they are, so I’m not sure why we had to do finite distances in the first place). To do this, I’ll introduce a rearrangement of \(\nabla_{\theta} \pi_\theta (s, a)\) and a score function:

\[\begin{align*} \nabla_\theta \pi_\theta (s, a) &= \pi_\theta (s, a) \frac{\nabla_\theta \pi_\theta (s, a)}{\pi_\theta (s, a)} \\ &= \pi_\theta (s, a) \underbrace{\nabla_\theta \log \pi_\theta (s, a)}_{\textstyle \text{score function}} \end{align*}\]

This formulation is convenient for us because the RHS has a \(\pi_\theta (s, a)\) factor that helps us take expectations of \(\nabla_\theta \pi_\theta (s, a)\).

Example: One-Step MDPs

Let’s see the above convenience at play in the case of a one-step MDP, starting in state \(s \sim d(s)\). Suppose the objective is this:

\[\begin{align*} J(\theta) = \mathbb{E}_{\pi_\theta} [r] \end{align*}\]

Then plugging in the score function reformulation:

\[\begin{align*} J(\theta) &= \sum_{s \in \mathcal{S}} d(s) \sum_{a \in \mathcal{A}} \pi_\theta \left( s, a \right) \mathcal{R}_{s, a} \\ \Rightarrow \nabla_\theta J(\theta) &= \sum_{s \in \mathcal{S}} d(s) \sum_{a \in \mathcal{A}} \pi_\theta \left( s, a \right) \nabla_\theta \log \pi_\theta (s, a) \mathcal{R}_{s, a} \\ &= \mathbb{E}_{\pi_\theta} [ \nabla_\theta \log \pi_\theta (s, a) r] \end{align*}\]

7.4. Policy Gradient Theorem

The policy gradient theorem simply states that for any of the earlier Policy Objective Functions we mentioned, the above equation for the policy objective function holds true, just that we use $Q$ instead of $r$:

\[\begin{align*} \nabla_\theta J(\theta) &= \mathbb{E}_{\pi_\theta} [ \nabla_\theta \log \pi_\theta (s, a) Q_{\pi_\theta}(s, a)] \end{align*}\]

NOTE: This is something I would just memorize, without remembering the (even if just intuitive) derivation

Advantage Function

In fact, just like how we used $Q$ instead of $r$, we can also just plug in $A$, as in the Advantage Function:

\[\begin{align*} A_\pi (s, a) &= Q_\pi(s, a) - V_\pi (s) \end{align*}\]

Alright, so we have $V$, $Q$, and now we have $A$. They all look like variants of the same thing, and indeed, in the end, the learnt policy (assuming some form of greedy) is still the same. But the advantage function is useful because it has lower variance.

7.5. Implementation: Monte-Carlo Policy Gradient (REINFORCE)

In the above equation, you’ll notice that there’s a \(\mathbb{E}_{\pi_\theta}\). But as with every iterative algorithm, we don’t actually have to compute that. We just naturally keep an estimate of it through sampling. Here’s a rough algorithm:

def REINFORCE(policy, learning_rate):
    policy.theta.init()
    for episode in episode_bank:
        for t in range(0, len(episode)):
            s, a, G_t = get_at_timestep(t, episode)
            loss = -learning_rate * torch.log(policy(s, a)) * G_t
            loss.backward()

This algorithm has a number of names - Monte-Carlo Policy Gradient, REINFORCE, Vanilla Policy Gradient, etc..

7.6. Actor-Critic Methods

REINFORCE is a Monte-Carlo method, meaning that we sampled full trajectories and used ground-truth returns \(G_t\) as our value function. We can also turn this into a shallow version (i.e. introduce bootstrapping) by introducing our friendly old learnt \(Q_w (s,a)\). Note that we denote $Q$’s weights as $w$ and $\pi$’s weights as $\theta$. Then, our policy gradient becomes:

\[\begin{align*} \mathbb{E} \left[ \nabla_\theta \log \pi_\theta (s, a) Q_w (s, a) \right] \end{align*}\]

Implementation

This means that we can apply pretty much the same algorithms (SARSA / DQN; they’re all the same save for some tricks here and there):

def Q_action_critic(env, policy, Q, policy_lr, Q_lr):
    policy.theta.init()
    Q.w.init()
    a = policy.sample_action(env.s)

    for _ in range(however_many_steps):
        s_next, r = env.step(a)
        a_next = policy.sample_action(s_next)

        TD_error = r + Q_lr * Q(s_next, a_next) - Q(s, a)
        policy_grad = -policy_lr * torch.log(policy(s, a)) * Q(s, a)

        TD_error.backward()
        policy_grad.backward()

        a = a_next
        s = s_next

You can use all of the same tricks in the $Q$ part of this implementation (e.g. use 2 $Q$ networks, use soft-updates, etc.).

8. Beyond Vanilla Policy Gradient

Since I’d like for us to achieve an understanding of RL sufficiently to extend it to LLM training, where PPO and GRPO are common, I will cover them here.

8.1. Proximal Policy Optimization

A 2017 method developed by John Schulman based on his earlier algorithm, Trust Region Policy Optimzation (TRPO). Both PPO and TRPO are motivated by the same question - how big of an update step can you take to improve performance without overstepping and causing performance collapse? For that, we have to introduce a new objective function:

\[\begin{align*} L (s, a, \theta_k, \theta) &= \min \left( \frac{\pi_\theta (a \mid s)}{\pi_{\theta_k} (a \mid s)} A_{\pi_{\theta_k}} (s, a), \text{clip}\left( \frac{\pi_\theta (a \mid s)}{\pi_{\theta_k} (a \mid s)} A_{\pi_{\theta_k}}, 1 - \epsilon, 1 + \epsilon \right) A_{\pi_{\theta_k}} (s, a) \right) \end{align*}\]

where $\theta_k$ is the old policy and $\theta$ is the updated policy.