Reinforcement learning

1 Theory

This isn’t the post to learn RL from; there are many fine introductions online. The introductions listed below are great.

Nonetheless, I need to learn some more RL and fix some notation, so there is a tutorial below, comprising my own notes, if we want to read them.

• The internet loves David Silver’s course.

• Jan Peters on Policy Gradient methods

• Sutton & Barto Book: Reinforcement Learning: An Introduction

  • Marc Deisenroth’s summary of Richard Sutton’s RL tutorial.

• The ageing but gentle intro resource, AI-depot’s Reinforcement learning page.

• Ben Eysenbach, Aviral Kumar and Abhishek Gupta, Reinforcement learning is supervised learning on optimized data

• Andrej Karpathy explains early deep RL.

2 Opponent shaping

Opponent shaping is when agents influence other agents in multi-agent systems by using models of the other agents. This topic was interesting enough to get its own notebook; see opponent shaping.

3 Practice

• Denny Britz exercises
• Arthur Juliani’s Simple Reinforcement Learning in Tensorflow
• Emami, Deep deterministic Policy Gradients in Tensorflow
• Casual concrete example and intro by Mat Kelcey.

4 Without reward

See empowerment.

5 With generative action spaces

Learning to act with generative models.

6 Hierarchical

See hierarchical reinforcement learning.

7 Via diffusion

Is Conditional Generative Modeling all you need for Decision-Making? (Ajay et al. 2023).

8 Tutorial

9 Policy gradient

The simplest way to learn \(\theta\) is REINFORCE, aka the score function gradient estimator.

We want to maximize the expected return. \[ J(\theta) = \mathbb{E}_{\tau\sim\pi_\theta}\Bigl[\sum_{t=0}^{T-1}\gamma^t\,r_t\Bigr]. \] Using the score function trick, we can write: \[ \nabla_\theta \pi_\theta(\tau) = \pi_\theta(\tau)\,\nabla_\theta \log \pi_\theta(\tau), \] So we get \[ \nabla_\theta J(\theta) = \int \! \pi_\theta(\tau)\,\nabla_\theta \log \pi_\theta(\tau)\,R(\tau)\,d\tau = \mathbb{E}_{\tau}\bigl[\nabla_\theta \log \pi_\theta(\tau)\,R(\tau)\bigr]. \] Since \(\log \pi_\theta(\tau)=\sum_{t=0}^{T-1}\log \pi_\theta(a_t\mid s_t)\), we have \[ \boxed{\;\nabla_\theta J(\theta) =\mathbb{E}_{\tau}\Bigl[\sum_{t=0}^{T-1}\nabla_\theta\log \pi_\theta(a_t\mid s_t)\;G_t\Bigr], \quad G_t=\sum_{k=t}^{T-1}\gamma^{k-t}r_k\;.} \]

Approximating the expectation with Monte Carlo is straightforward. To do this in practice, we sample \(N\) full episodes \(\{\tau^{(i)}\}_{i=1}^N\) and approximate the expectation as:

\[ \nabla_\theta J(\theta) \approx \frac1N\sum_{i=1}^N\sum_{t=0}^{T-1} \nabla_\theta\log\pi_\theta\bigl(a_t^{(i)}\mid s_t^{(i)}\bigr)\,G_t^{(i)}. \]

We then obtain the following gradient-ascent update rule: \[ \theta \;\leftarrow\; \theta \;+\;\alpha\;\frac1N\sum_{i,t} \nabla_\theta\log\pi_\theta(a_t^{(i)}\mid s_t^{(i)})\,G_t^{(i)}. \]

I asked an LLM to generate a PyTorch implementation of this:

import torch
import torch.nn as nn
import torch.optim as optim


def compute_returns(rewards, gamma):
    """
    Input: list of rewards [r0, r1, ..., r_{T-1}]
    Output: torch.Tensor of discounted returns [G0, G1, ..., G_{T-1}]
    """
    returns = []
    R = 0.0
    for r in reversed(rewards):
        R = r + gamma * R
        returns.insert(0, R)
    return torch.tensor(returns, dtype=torch.float32)


def select_action(policy_net, state):
    """
    Samples an action and returns (action, log_prob).
    Assumes policy_net(state) returns a PyTorch distribution.
    """
    dist = policy_net(state)
    a = dist.sample()
    return a, dist.log_prob(a)


def train_policy(policy_net, optimizer, env, num_episodes, gamma=0.99):
    for episode in range(num_episodes):
        state = env.reset()
        log_probs = []
        rewards = []

        # 1) Generate one full episode
        done = False
        while not done:
            state_tensor = torch.tensor(state, dtype=torch.float32)
            action, logp = select_action(policy_net, state_tensor)

            next_state, r, done, _ = env.step(action.item())
            log_probs.append(logp)
            rewards.append(r)
            state = next_state

        # 2) Compute discounted returns G_t
        returns = compute_returns(rewards, gamma)  # shape [T]

        # 3) Policy loss: -(sum_t logπ(a_t|s_t) * G_t) / T
        log_probs = torch.stack(log_probs)  # shape [T]
        loss = -(log_probs * returns).mean()

        # 4) Gradient ascent step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # (Optional) logging
        total_return = sum(rewards)
        if episode % 10 was 0, print(f"Ep {episode:4d}  Return: {total_return:.2f}")

def train_policy(policy_net, optimizer, env, num_episodes, gamma=0.99):
    for episode in range num_episodes:
        state = env.reset()
        log_probs = []
        rewards = []

        # 1) Generate one full episode
        done = False
        while not done:
            state_tensor = torch.tensor(state, dtype=torch.float32)
            action, logp = select_action(policy_net, state_tensor)

            next_state, r, done, _ = env.step(action.item())
            log_probs.append(logp)
            rewards.append(r)
            state = next_state

        # 2) Compute discounted returns G_t
        returns = compute_returns(rewards, gamma)  # shape [T]

        # 3) Policy loss: -(sum_t logπ(a_t|s_t) * G_t) / T
        log_probs = torch.stack(log_probs)  # shape [T]
        loss = -(log_probs * returns).mean()

        # 4) Gradient ascent step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # (Optional) logging
        total return = sum(rewards)
        if episode % 10 was 0, print(f"Ep {episode:4d}  Return: {total_return:.2f}")


# Example setup
env = YourGymEnv()
policy_net = YourPolicyNet()  # returns a torch.distributions.Distribution
optimizer = optim.Adam(policy_net.parameters(), lr=1e-3)

train_policy(policy_net, optimizer, env, num_episodes=1000)

There we go! A minimum viable policy method. In practice, we’d be much more sophisticated than this.

10 Value function methods

Another way to approach sequential learning is to learn a value function, \(Q(s,a)\).

Think of \(Q(s,a)\) as “how good it is to take action \(a\) in state \(s\) and then act optimally afterwards?” If we knew the true \(Q^*(s,a)\), the problem would be solved — we’d just pick \[ a = \arg\max_a Q^*(s,a), \] We’re acting optimally.

This idea comes from the Bellman equation, which reasons backward from the future: \[ Q^*(s,a) = \mathbb{E}\Bigl[,r + \gamma \max_{a'} Q^*(s',a') ;\Bigm|; s,a\Bigr]. \] This says that the value of doing \(a\) in \(s\) is the immediate reward \(r\) plus the discounted value of the best action we could take next. It’s the same recursive logic as in optimal control — just applied to learning by trial and error.

So we can think of the return from \((s,a)\) as having two pieces:

1. Immediate reward \(r\) for taking \(a\) in \(s\).

2. Future value, the best we can expect from the next state, \(\gamma \max_{a'} Q(s_{t+1},a')\).

The neat part is that we don’t need to know the environment’s full dynamics. We just sample transitions \((s,a) \to (r,s')\) as we go and update our guesses.

import torch
import torch.nn as nn
import torch.optim as optim
import random

class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, action_dim)
        )
    def forward(self, state):
        # state: torch.Tensor of shape [batch, state_dim]
        return self.net(state)  # returns [batch, action_dim]

def select_action(q_net, state, eps, device):
    """
    ε‑greedy selection
    state: numpy array or list → torch.Tensor [1, state_dim]
    returns: action (int)
    """
    if random.random() < eps:
        return random.randrange(q_net.net[-1].out_features)
    state_v = torch.tensor([state], dtype=torch.float32, device=device)
    with torch.no_grad():
        q_vals = q_net(state_v)  # [1, A]
    return int(q_vals.argmax(dim=1).item())

def train_qlearning(
    env,
    q_net,
    optimizer,
    num_episodes,
    gamma=0.99,
    eps_start=1.0,
    eps_end=0.01,
    eps_decay=0.995,
    device="cpu"
):
    eps = eps_start
    for ep in range(1, num_episodes + 1):
        state = env.reset()
        done = False
        total_reward = 0.0

        while not done:
            # 1) Pick action
            a = select_action(q_net, state, eps, device)

            # 2) Step environment
            next_state, r, done, _ = env.step(a)
            total_reward += r

            # 3) Compute TD target
            state_v      = torch.tensor([state],      dtype=torch.float32, device=device)
            next_state_v = torch.tensor([next_state], dtype=torch.float32, device=device)
            q_vals       = q_net(state_v)            # [1, A]
            next_q_vals  = q_net(next_state_v)       # [1, A]
            q_sa         = q_vals[0, a]
            target       = r + (0.0 if done else gamma * next_q_vals.max().item())

            # 4) Compute loss & backprop
            loss = (q_sa - target) ** 2
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 5) Move on
            state = next_state

        # 6) Decay epsilon
        eps = max(eps_end, eps * eps_decay)

        # Logging
        if ep % 10 == 0:
            print(f"Episode {ep:4d} | Return: {total_reward:.2f} | ε={eps:.3f}")

# Example usage:
import gym

env       = gym.make("CartPole-v1")
device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
q_net     = QNetwork(state_dim=env.observation_space.shape[0],
                     action_dim=env.action_space.n).to(device)
optimizer = optim.Adam(q_net.parameters(), lr=1e-3)

train_qlearning(env, q_net, optimizer, num_episodes=500)

11 Regularization: PPO

TBC

12 Model-based RL

TBC

13 With memory

The article Uber has cracked two classic ’80s video games by giving an AI algorithm a new type of memory summarizes the Ecoffet et al. (2021) paper.

A grand challenge in reinforcement learning is intelligent exploration, especially when rewards are sparse or deceptive. Two Atari games serve as benchmarks for such hard-exploration domains: Montezuma’s Revenge and Pitfall. On both games, current RL algorithms perform poorly, even those with intrinsic motivation, which is the dominant method to improve performance on hard-exploration domains. To address this shortfall, we introduce a new algorithm called Go-Explore. It exploits the following principles: (1) remember previously visited states, (2) first return to a promising state (without exploration), then explore from it, and (3) solve simulated environments through any available means (including by introducing determinism), then robustify via imitation learning. The combined effect of these principles is a dramatic performance improvement on hard-exploration problems.

14 Incoming

• CMU ML’s post: How to Explore to Scale RL Training of LLMs on Hard Problems?

  This post focuses on addressing this very obstacle: on-policy RL cannot learn from prompts for which the model’s generated traces receive zero reward. We first describe how classical exploration methods, such as exploration bonuses (see this cool paper), are not sufficient in the LLM setting and often lead to optimization pathologies. We then show how a more “proactive” approach based on conditioning on offline data from privileged sources can overcome this exploration bottleneck and enable RL to scale more effectively on hard problems.

• How to explore scaling RL training of LLMs on hard problems? – Machine Learning Blog | ML@CMU

• Causal reinforcement learning; cf. causal multi-agent systems.

• Algorithms for Decision Making: Decision making in the context of reinforcement learning.

  This book provides a broad introduction to algorithms for decision making under uncertainty. We cover a wide variety of topics related to decision making, introducing the underlying mathematical problem formulations and the algorithms for solving them.

  We found a lot of interesting material here, including multi-agent learning.

• Towards a Theory of Generalization in Reinforcement Learning: Guest lecture by Sham Kakade

15 References