45.3 Reinforcement Learning as Optimization

1. Discovery as a Markov Decision Process

A Markov Decision Process (MDP) is defined by the tuple \((\mathcal{S}, \mathcal{A}, P, R, \gamma)\):

An MDP models any situation where an agent makes sequential decisions under uncertainty, receiving feedback after each step. It matters because it transforms vague notions of "planning" into a precise optimization problem with provable solution methods. Given the state, transition, and reward structure, algorithms compute the decision rule (policy) that maximizes long-term cumulative reward. The core mechanism is the Bellman equation, which splits a multi-step plan's value into the immediate reward plus the discounted value of the next state. This decomposition enables dynamic programming or gradient-based solutions. Use an MDP formulation (and therefore RL) instead of single-shot optimization when earlier decisions affect later options. If each evaluation is independent, simpler methods like Bayesian optimization or evolutionary search are more sample-efficient.

• States \(\mathcal{S}\): the current state of knowledge (observed data, remaining budget, intermediate compounds synthesized so far).
• Actions \(\mathcal{A}\): the set of possible next experiments, reactions, or measurements.
• Transition dynamics \(P(s' \mid s, a)\): how the state evolves after taking action \(a\) in state \(s\). In discovery, this includes both the deterministic part (budget decreases, data accumulates) and the stochastic part (experiment outcomes are noisy).
• Reward \(R(s, a, s')\): the immediate value of the transition. This might be the activity of a newly synthesized compound, the information gained from a measurement, or a penalty for resource consumption.
• Discount factor \(\gamma \in [0, 1]\): how much the agent values future rewards relative to immediate ones.

The goal is to find a policy \(\pi(a \mid s)\) (a distribution over actions given the current state) that maximizes the expected cumulative reward:

where \(\tau = (s_0, a_0, s_1, a_1, \ldots)\) is a trajectory sampled by following policy \(\pi\) in the MDP.

Mapping a discovery campaign to an MDP requires careful state design. Table 45.1 shows three example mappings.

2. The Policy Gradient Theorem

The MDP formulation raises a core question: how to find the policy that maximizes cumulative reward.

If the agent cannot estimate how a small change to its decision rule will affect cumulative reward, it is reduced to random trial and error, wasting every experiment in a budget-limited campaign. The policy gradient theorem solves exactly this problem, providing an unbiased gradient signal that turns sequential discovery into tractable optimization.

We parameterize the policy as \(\pi_\theta(a \mid s)\) using a neural network with parameters \(\theta\). The policy gradient theorem (Sutton et al., 1999) gives us the gradient of the expected return with respect to \(\theta\):

where \(G_t = \sum_{k=t}^{T} \gamma^{k-t} R_k\) is the return from time step \(t\). This expression uses the log-derivative trick: because \(\nabla_\theta \pi_\theta = \pi_\theta \nabla_\theta \log \pi_\theta\), we can estimate the gradient by sampling trajectories under \(\pi_\theta\) and weighting each action's log-probability by the return that followed it. Actions that led to high returns get reinforced (their probability increases); actions that led to low returns get suppressed. In short: weight each decision by the outcome it produced, then nudge the policy toward decisions that paid off.

2.1 REINFORCE

The simplest policy gradient algorithm, REINFORCE (Williams, 1992), directly applies the policy gradient theorem using Monte Carlo estimates (where the return \(G_t\) is computed from a complete sampled trajectory rather than bootstrapped from a value function) of \(G_t\):

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np

class PolicyNetwork(nn.Module):
    """Simple policy network for discrete action spaces."""

    def __init__(self, state_dim: int, n_actions: int, hidden: int = 64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, n_actions),
        )

    def forward(self, state: torch.Tensor) -> Categorical:
        logits = self.net(state)
        return Categorical(logits=logits)


def reinforce(
    env,
    policy: PolicyNetwork,
    n_episodes: int = 1000,
    lr: float = 1e-3,
    gamma: float = 0.99,
) -> list:
    """Train a policy using REINFORCE with baseline subtraction.

    Parameters
    ----------
    env : environment with reset() -> state, step(action) -> (state, reward, done)
    policy : PolicyNetwork
    n_episodes : int
    lr : float
    gamma : float

    Returns
    -------
    episode_returns : list of float
        Total return per episode.
    """
    optimizer = optim.Adam(policy.parameters(), lr=lr)
    baseline = 0.0  # Running average baseline
    episode_returns = []

    for episode in range(n_episodes):
        state = env.reset()
        log_probs = []
        rewards = []

        # Collect a trajectory
        done = False
        while not done:
            state_t = torch.FloatTensor(state)
            dist = policy(state_t)
            action = dist.sample()
            log_probs.append(dist.log_prob(action))

            state, reward, done, _ = env.step(action.item())
            rewards.append(reward)

        # Compute discounted returns
        returns = []
        G = 0.0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        returns = torch.FloatTensor(returns)

        # Baseline subtraction reduces variance
        episode_return = returns[0].item()
        baseline = 0.99 * baseline + 0.01 * episode_return
        returns = returns - baseline

        # Policy gradient update
        loss = -torch.stack([lp * G for lp, G in zip(log_probs, returns)]).sum()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        episode_returns.append(episode_return)

    return episode_returns

REINFORCE is simple but suffers from high variance: each gradient estimate depends on the return of an entire trajectory, and trajectories vary widely. Two techniques reduce this variance: baseline subtraction (implemented above) and using a learned value function as the baseline, which leads us to actor-critic methods.

3. Actor-Critic Methods

An actor-critic architecture maintains two networks: the actor \(\pi_\theta(a \mid s)\) (the policy) and the critic \(V_\phi(s)\) (a value function that estimates the expected return from state \(s\)). The critic replaces the Monte Carlo return \(G_t\) with a lower-variance estimate, the advantage. Figure 45.3.1 illustrates Actor-Critic Architecture for Sequential Discovery.

The advantage measures how much better action \(a_t\) was compared to the average action from state \(s_t\). The policy gradient becomes:

Figure 45.3 illustrates how the actor and critic interact during training. The shared feature extractor feeds both heads, and the environment closes the loop by returning rewards and next states.

The critic is trained to minimize the squared temporal-difference (TD) error (the gap between the critic's current value estimate and a one-step-ahead bootstrap target, used as a training signal to improve the critic): \(\mathcal{L}_\text{critic} = \mathbb{E}[(R_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t))^2]\).

class ActorCritic(nn.Module):
    """Combined actor-critic network with shared feature extractor."""

    def __init__(self, state_dim: int, n_actions: int, hidden: int = 128):
        super().__init__()
        # Shared feature layers
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        # Actor head: policy logits
        self.actor = nn.Linear(hidden, n_actions)
        # Critic head: state value
        self.critic = nn.Linear(hidden, 1)

    def forward(self, state: torch.Tensor):
        features = self.shared(state)
        policy = Categorical(logits=self.actor(features))
        value = self.critic(features).squeeze(-1)
        return policy, value


def train_actor_critic(
    env,
    model: ActorCritic,
    n_episodes: int = 2000,
    lr: float = 3e-4,
    gamma: float = 0.99,
    entropy_coef: float = 0.01,
) -> list:
    """Train an actor-critic agent on a discovery environment.

    Parameters
    ----------
    env : environment with reset/step interface
    model : ActorCritic network
    n_episodes : int
    lr : float
    gamma : float
    entropy_coef : float
        Entropy bonus to encourage exploration.

    Returns
    -------
    episode_returns : list of float
    """
    optimizer = optim.Adam(model.parameters(), lr=lr)
    episode_returns = []

    for episode in range(n_episodes):
        state = env.reset()
        total_reward = 0.0
        done = False

        log_probs, values, rewards, entropies = [], [], [], []

        while not done:
            state_t = torch.FloatTensor(state)
            policy, value = model(state_t)
            action = policy.sample()

            log_probs.append(policy.log_prob(action))
            values.append(value)
            entropies.append(policy.entropy())

            state, reward, done, _ = env.step(action.item())
            rewards.append(reward)
            total_reward += reward

        # Compute advantages using TD residuals
        returns = []
        G = 0.0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        returns = torch.FloatTensor(returns)
        values_t = torch.stack(values)

        advantages = returns - values_t.detach()

        # Combined loss: actor + critic + entropy
        actor_loss = -(torch.stack(log_probs) * advantages).sum()
        critic_loss = 0.5 * ((returns - values_t) ** 2).sum()
        entropy_loss = -torch.stack(entropies).sum()

        loss = actor_loss + critic_loss + entropy_coef * entropy_loss

        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
        optimizer.step()

        episode_returns.append(total_reward)

    return episode_returns

4. Continuous Action Spaces

The actor-critic framework above assumes a finite set of discrete actions, but many discovery problems require choosing from a continuum of possibilities.

When actions are continuous (e.g., choosing a composition vector or setting continuous process parameters), the policy outputs the parameters of a continuous distribution, typically a diagonal Gaussian (a multivariate normal distribution where each dimension is independent, so the covariance matrix contains only diagonal entries):

The log-probability under this distribution is:

Gradient interpretation

The policy gradient then pushes the mean \(\mu_\theta\) toward actions that yielded high advantage and adjusts the variance \(\sigma_\theta\) to explore more (if exploration helps) or less (if exploitation is sufficient).

class ContinuousActorCritic(nn.Module):
    """Actor-critic for continuous action spaces (Gaussian policy)."""

    def __init__(self, state_dim: int, action_dim: int, hidden: int = 128):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.Tanh(),
            nn.Linear(hidden, hidden),
            nn.Tanh(),
        )
        self.mu_head = nn.Linear(hidden, action_dim)
        self.log_std = nn.Parameter(torch.zeros(action_dim))  # Learnable
        self.value_head = nn.Linear(hidden, 1)

    def forward(self, state: torch.Tensor):
        features = self.shared(state)
        mu = self.mu_head(features)
        std = self.log_std.exp().expand_as(mu)
        dist = torch.distributions.Normal(mu, std)
        value = self.value_head(features).squeeze(-1)
        return dist, value

    def act(self, state: torch.Tensor):
        """Sample an action and return it with log-probability and value."""
        dist, value = self.forward(state)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum(-1)  # Sum over action dimensions
        return action, log_prob, value

5. When RL Outperforms Static Optimization

RL is not always the right tool. Its advantages over BO and evolutionary methods emerge in specific settings:

• Long horizons with intermediate feedback: when the discovery process has many stages and each stage provides partial information (e.g., multi-step synthesis, sequential screening), RL learns to use that feedback adaptively.
• Shared structure across campaigns: if you run many similar discovery campaigns (e.g., optimizing different proteins with similar fitness landscapes), RL can meta-learn a policy that transfers across campaigns. BO starts from scratch each time.
• Resource allocation: when the decision is not "where to evaluate" but "how to allocate a fixed budget across competing subproblems," RL naturally handles the sequential budget constraint.
• Combinatorial action spaces: reaction planning, molecule assembly, and protocol design involve choosing from structured combinatorial spaces where BO's continuous surrogates are a poor fit.

RL struggles when:

• The evaluation budget is tiny (\(N < 50\)): RL needs many episodes to learn a policy. With very few evaluations, BO's sample efficiency wins.
• The environment dynamics are unknown and hard to simulate: RL needs either a simulator or many real interactions. If neither is available, BO with a Gaussian process (GP) surrogate is more data-efficient.
• The problem is single-step: if there is no sequential structure, RL reduces to bandit optimization (selecting the best option from a set without sequential state transitions), where BO's surrogate model typically provides better sample efficiency than policy gradient methods.

6. Scaling with Ray Tune

Knowing when RL fits is only half the battle; training a policy to convergence can itself become a computational bottleneck.

Training RL policies requires many episodes, and each episode may involve expensive evaluations. Ray Tune distributes this workload across a cluster:

from ray import tune
from ray.tune.schedulers import ASHAScheduler

def train_rl_agent(config):
    """Training function compatible with Ray Tune."""
    env = DiscoveryEnv(
        n_compounds=config["n_compounds"],
        noise_level=config["noise_level"],
    )
    model = ActorCritic(
        state_dim=env.observation_space.shape[0],
        n_actions=env.action_space.n,
        hidden=config["hidden_size"],
    )
    returns = train_actor_critic(
        env, model,
        n_episodes=config["n_episodes"],
        lr=config["lr"],
        gamma=config["gamma"],
        entropy_coef=config["entropy_coef"],
    )
    # Report final performance to Ray Tune
    tune.report(mean_return=np.mean(returns[-100:]))


# Hyperparameter search over RL training configs
scheduler = ASHAScheduler(
    metric="mean_return",
    mode="max",
    max_t=5000,
    grace_period=500,
    reduction_factor=3,
)

analysis = tune.run(
    train_rl_agent,
    config={
        "n_compounds": 100,
        "noise_level": 0.1,
        "hidden_size": tune.choice([64, 128, 256]),
        "lr": tune.loguniform(1e-4, 1e-2),
        "gamma": tune.uniform(0.95, 0.999),
        "entropy_coef": tune.loguniform(1e-3, 1e-1),
        "n_episodes": 5000,
    },
    num_samples=20,
    scheduler=scheduler,
)

from stable_baselines3 import PPO

model = PPO("MlpPolicy", env, verbose=1, learning_rate=3e-4)
model.learn(total_timesteps=100_000)

# Deploy the trained policy
obs = env.reset()
for _ in range(50):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, info = env.step(action)

7. Multi-Objective RL

When the discovery MDP has multiple reward signals (e.g., compound activity and synthetic cost at each step), the policy gradient extends to handle vector-valued rewards. The simplest approach applies linear scalarization (combining multiple objectives into a single number by weighting each one): \(R = \sum_i w_i R_i\) with weights \(\mathbf{w}\) encoding the scientist's preferences. A more sophisticated approach trains a family of policies indexed by \(\mathbf{w}\), producing a Pareto front (the set of solutions where no objective can be improved without worsening another) of policies analogous to the Pareto front of solutions in Non-dominated Sorting Genetic Algorithm III (NSGA-III). Recent multi-objective RL (MORL) methods maintain a set of policies that span the achievable trade-offs, connecting RL to the multi-objective framework of Section 45.2.

The Discovery Workbench integrates RL agents through the same suggest/observe interface as the Bayesian optimizer in Section 45.1. The key deployment difference: RL agents maintain internal state (policy activations or episode context) and require checkpointing, while BO agents are stateless, reconstructing their strategy from the observation history alone.