Chapter 1: Why Reinforcement Learning for Materials Science

Grasp how to optimize materials and processes within a sequential decision-making framework. We also introduce the pitfalls of reward design.

💡 Note: A reward is a "prize for an action." If you get the balance between short-term and long-term prizes wrong, learning goes astray.

Learning Objectives

In this chapter, you will learn the following:

• The limitations of conventional methods in materials discovery and the role of reinforcement learning
• Basic concepts of the Markov Decision Process (MDP)
• How Q-learning and the Deep Q-Network (DQN) work
• Implementation for a simple materials discovery task

1.1 Challenges in Materials Discovery and the Role of Reinforcement Learning

Limitations of Conventional Materials Discovery

Developing new materials involves an enormous search space (composition, structure, process conditions):

• Composition search: Just choosing 3 elements from the periodic table gives $\binom{118}{3} \approx 267,000$ combinations
• Structure search: There are 230 space groups for crystal structures alone
• Process search: The combinations of temperature, pressure, and time are infinite

With the conventional trial-and-error approach: - Reliance on the researcher's experience and intuition - Time- and cost-intensive evaluation (several weeks to several months per material) - Prone to getting stuck in local optima

Problems: 1. Inefficient: Repeatedly tries similar materials 2. Narrow search: Limited to the scope of the researcher's knowledge 3. Low reproducibility: Reliant on tacit knowledge

Solution via Reinforcement Learning

Reinforcement learning is a framework that learns optimal actions through interaction with an environment:

Advantages of reinforcement learning: 1. Automatic optimization: Automates trial and error and learns efficient search strategies 2. Balance of exploration and exploitation: Adjusts between exploring unknown regions and exploiting known good regions 3. Sequential improvement: Learns from each evaluation result and improves the next choice 4. Closed loop: Can integrate with experimental equipment and operate 24 hours a day

Success Stories in Materials Science

Example 1: Optimization of Li-ion battery electrolytes (MIT, 2022) - Challenge: Optimize the blending ratios of 5 components (search space > $10^6$) - Method: Sequentially select formulations with DQN - Result: Found the optimal solution 5 times faster than conventional methods, with a 30% improvement in ionic conductivity

Example 2: Donor materials for organic solar cells (University of Toronto, 2021) - Challenge: Optimize molecular structure (10^23 candidates) - Method: Integrate molecular generation and evaluation with Actor-Critic - Result: Discovered a new material with 15% photoelectric conversion efficiency in 3 months (conventionally 2 years)

1.2 Fundamentals of the Markov Decision Process (MDP)

What Is an MDP?

The mathematical foundation of reinforcement learning is the Markov Decision Process (MDP). An MDP is defined by the following 5-tuple:

$$ \text{MDP} = (S, A, P, R, \gamma) $$

• $S$: State space (e.g., properties of currently tried materials)
• $A$: Action space (e.g., the next material candidate to try)
• $P(s'|s, a)$: State transition probability (the probability of transitioning from state $s$ to $s'$ when taking action $a$)
• $R(s, a, s')$: Reward function (the reward obtained from a state transition)
• $\gamma \in [0, 1)$: Discount factor (the importance of future rewards)

Mapping to Materials Discovery

The Markov Property

An important assumption of the MDP is the Markov property:

$$ P(s_{t+1}|s_t, a_t, s_{t-1}, a_{t-1}, \dots) = P(s_{t+1}|s_t, a_t) $$

In other words, the next state depends only on the current state and action, and the past history is unnecessary.

In materials discovery, if you choose the next material (action) based on the current evaluation results (state), there is no need to remember the entire past history.

Policy and Value Functions

Policy $\pi(a|s)$: The probability of choosing action $a$ in state $s$

State value function $V^\pi(s)$: The expected cumulative reward when acting according to policy $\pi$ starting from state $s$

$$ V^\pi(s) = \mathbb{E}_\pi \left[ \sum_{t=0}^\infty \gamma^t r_t \mid s_0 = s \right] $$

Action value function (Q-function) $Q^\pi(s, a)$: The expected cumulative reward when taking action $a$ in state $s$ and then following policy $\pi$

$$ Q^\pi(s, a) = \mathbb{E}_\pi \left[ \sum_{t=0}^\infty \gamma^t r_t \mid s_0 = s, a_0 = a \right] $$

Optimal policy $\pi^*$: The policy that maximizes the value function in every state

$$ \pi^* = \arg\max_\pi V^\pi(s) \quad \forall s \in S $$

1.3 Q-Learning

The Basic Idea of Q-Learning

Q-learning is a reinforcement learning algorithm that learns the Q-function directly.

Bellman equation: $$ Q^*(s, a) = \mathbb{E}_{s'} \left[ r + \gamma \max_{a'} Q^*(s', a') \mid s, a \right] $$

This means "the optimal Q-function equals the discounted sum of the immediate reward $r$ and the maximum Q-value in the next state."

The Q-Learning Update Rule

Based on the observed transition $(s, a, r, s')$, update the Q-value:

$$ Q(s, a) \leftarrow Q(s, a) + \alpha \left[ r + \gamma \max_{a'} Q(s', a') - Q(s, a) \right] $$

• $\alpha$: Learning rate (0 to 1)
• $r + \gamma \max_{a'} Q(s', a')$: TD target (Temporal Difference Target)
• $r + \gamma \max_{a'} Q(s', a') - Q(s, a)$: TD error

Implementation in Python

We implement Q-learning in a simple grid world (a metaphor for the materials search space):

"""
Implementation of a materials discovery environment using Q-learning

Dependencies (libraries and versions):
- Python: 3.9+
- numpy: 1.24+
- matplotlib: 3.7+

Reproducibility:
- Fixed random seed: 42 (unified across all random operations)
- Number of episodes: 1000 (convergence confirmed)
- Learning rate alpha: 0.1 (prevents excessive updates)
- Discount factor gamma: 0.99 (emphasizes long-term rewards)
- epsilon-greedy: epsilon=0.1 fixed (10% exploration, 90% exploitation)

Pitfalls (practical caveats):
1. Because epsilon is fixed, exploration continues even late in training (room for optimization)
2. Q-table size is 5x5x4=100 elements (only suitable for small environments)
3. Rewards are sparse (+10 only at the goal), so exploration may be difficult
"""

import numpy as np
import matplotlib.pyplot as plt

# Fix the random seed (ensure reproducibility)
RANDOM_SEED = 42
np.random.seed(RANDOM_SEED)

class SimpleMaterialsEnv:
    """Simple materials discovery environment (grid world)

    - 5x5 grid
    - Each cell represents a material candidate
    - Goal: reach the material with the best properties (the goal)
    """
    def __init__(self):
        self.grid_size = 5
        self.state = (0, 0)  # Start position
        self.goal = (4, 4)   # Goal position (optimal material)

    def reset(self):
        """Reset to the initial state"""
        self.state = (0, 0)
        return self.state

    def step(self, action):
        """Execute an action

        Args:
            action: 0=up, 1=down, 2=left, 3=right

        Returns:
            next_state, reward, done
        """
        x, y = self.state

        # Move according to the action
        if action == 0 and x > 0:  # Up
            x -= 1
        elif action == 1 and x < self.grid_size - 1:  # Down
            x += 1
        elif action == 2 and y > 0:  # Left
            y -= 1
        elif action == 3 and y < self.grid_size - 1:  # Right
            y += 1

        self.state = (x, y)

        # Reward design
        if self.state == self.goal:
            reward = 10.0  # Reached the goal (optimal material found)
            done = True
        else:
            reward = -0.1  # Cost of each step (experimental cost)
            done = False

        return self.state, reward, done

    def get_state_space(self):
        """Size of the state space"""
        return self.grid_size * self.grid_size

    def get_action_space(self):
        """Size of the action space"""
        return 4


def q_learning(env, episodes=1000, alpha=0.1, gamma=0.99, epsilon=0.1):
    """Q-learning algorithm

    Args:
        env: environment
        episodes: number of episodes
        alpha: learning rate
        gamma: discount factor
        epsilon: probability of epsilon-greedy exploration

    Returns:
        The learned Q-table
    """
    # Initialize the Q-table (state x action)
    Q = np.zeros((env.grid_size, env.grid_size, env.get_action_space()))

    rewards_per_episode = []

    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        done = False

        while not done:
            # epsilon-greedy exploration
            if np.random.random() < epsilon:
                action = np.random.randint(env.get_action_space())  # Random exploration
            else:
                action = np.argmax(Q[state[0], state[1], :])  # Exploit the best action

            # Execute the action
            next_state, reward, done = env.step(action)
            total_reward += reward

            # Update the Q-value (Bellman equation)
            current_q = Q[state[0], state[1], action]
            max_next_q = np.max(Q[next_state[0], next_state[1], :])
            new_q = current_q + alpha * (reward + gamma * max_next_q - current_q)
            Q[state[0], state[1], action] = new_q

            state = next_state

        rewards_per_episode.append(total_reward)

        if (episode + 1) % 100 == 0:
            avg_reward = np.mean(rewards_per_episode[-100:])
            print(f"Episode {episode+1}: Avg Reward = {avg_reward:.2f}")

    return Q, rewards_per_episode


# Execution
env = SimpleMaterialsEnv()
Q, rewards = q_learning(env, episodes=1000)

# Visualize the learning curve
plt.figure(figsize=(10, 6))
plt.plot(np.convolve(rewards, np.ones(50)/50, mode='valid'))
plt.xlabel('Episode')
plt.ylabel('Average Reward (50 episodes)')
plt.title('Q-Learning: Learning Progress in the Materials Discovery Environment')
plt.grid(True)
plt.show()

# Visualize the learned Q-values
policy = np.argmax(Q, axis=2)
print("\nLearned policy (best action in each cell):")
print("0=up, 1=down, 2=left, 3=right")
print(policy)

Example output:

Episode 100: Avg Reward = -4.52
Episode 200: Avg Reward = -3.21
Episode 500: Avg Reward = -1.85
Episode 1000: Avg Reward = -1.12

Learned policy (best action in each cell):
[[1 1 1 1 1]
 [1 1 1 1 1]
 [1 1 1 1 1]
 [1 1 1 1 1]
 [3 3 3 3 0]]

Explanation: - Initially the reward is low (-4.52), but it improves as learning progresses (-1.12) - Eventually, the shortest path to the goal is learned (a down-then-right policy)

1.4 Deep Q-Network (DQN)

Limitations of Q-Learning

Q-learning is effective when states and actions are discrete and few in number. However, in materials science:

• The state space is huge: material descriptors (100+ dimensions)
• Continuous values: composition ratios, temperature, pressure, etc.
• The Q-table becomes impractical: you cannot store $10^{100}$ cells

The DQN Solution

DQN approximates the Q-function with a neural network:

$$ Q(s, a; \theta) \approx Q^*(s, a) $$

• $\theta$: parameters of the neural network

Loss function: $$ L(\theta) = \mathbb{E}_{(s,a,r,s') \sim D} \left[ \left( r + \gamma \max_{a'} Q(s', a'; \theta^-) - Q(s, a; \theta) \right)^2 \right] $$

• $D$: experience replay buffer (stores past transitions)
• $\theta^-$: target network (stabilizes learning)

Key Techniques of DQN

1. Experience Replay: Randomly samples past transitions to reduce data correlation
2. Target Network: Computes the TD target with a fixed network to stabilize learning
3. epsilon-greedy exploration: Balances exploration (random) and exploitation (best action)

DQN Implementation in PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque
import random

class DQN(nn.Module):
    """Deep Q-Network

    Takes a state as input and outputs the Q-value for each action
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)


class ReplayBuffer:
    """Experience replay buffer"""
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)

    def __len__(self):
        return len(self.buffer)


class DQNAgent:
    """DQN agent"""
    def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99, epsilon=1.0, epsilon_decay=0.995, epsilon_min=0.01):
        self.action_dim = action_dim
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min

        # Main network and target network
        self.policy_net = DQN(state_dim, action_dim)
        self.target_net = DQN(state_dim, action_dim)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=lr)
        self.buffer = ReplayBuffer()

    def select_action(self, state):
        """epsilon-greedy action selection"""
        if np.random.random() < self.epsilon:
            return np.random.randint(self.action_dim)
        else:
            with torch.no_grad():
                state_tensor = torch.FloatTensor(state).unsqueeze(0)
                q_values = self.policy_net(state_tensor)
                return q_values.argmax().item()

    def train(self, batch_size=64):
        """Mini-batch training"""
        if len(self.buffer) < batch_size:
            return

        # Mini-batch sampling
        batch = self.buffer.sample(batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions).unsqueeze(1)
        rewards = torch.FloatTensor(rewards).unsqueeze(1)
        next_states = torch.FloatTensor(next_states)
        dones = torch.FloatTensor(dones).unsqueeze(1)

        # Current Q-values
        current_q = self.policy_net(states).gather(1, actions)

        # Target Q-values
        with torch.no_grad():
            max_next_q = self.target_net(next_states).max(1)[0].unsqueeze(1)
            target_q = rewards + (1 - dones) * self.gamma * max_next_q

        # Loss computation and optimization
        loss = nn.MSELoss()(current_q, target_q)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Decay epsilon
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

    def update_target_network(self):
        """Update the target network"""
        self.target_net.load_state_dict(self.policy_net.state_dict())


# Materials discovery environment (continuous-state version)
class ContinuousMaterialsEnv:
    """Materials discovery environment with a continuous state space"""
    def __init__(self, state_dim=4):
        self.state_dim = state_dim
        self.target = np.array([3.0, 5.0, 2.5, 4.0])  # Target properties
        self.state = None

    def reset(self):
        self.state = np.random.uniform(0, 10, self.state_dim)
        return self.state

    def step(self, action):
        # Actions: 0=increase, 1=decrease, 2=large increase, 3=large decrease
        delta = [0.1, -0.1, 0.5, -0.5][action]

        # Change a random dimension
        dim = np.random.randint(self.state_dim)
        self.state[dim] = np.clip(self.state[dim] + delta, 0, 10)

        # Reward: distance from the target (negative; closer is better)
        distance = np.linalg.norm(self.state - self.target)
        reward = -distance

        # Termination condition: sufficiently close to the target
        done = distance < 0.5

        return self.state, reward, done


# DQN training
env = ContinuousMaterialsEnv()
agent = DQNAgent(state_dim=4, action_dim=4)

episodes = 500
rewards_history = []

for episode in range(episodes):
    state = env.reset()
    total_reward = 0
    done = False

    while not done:
        action = agent.select_action(state)
        next_state, reward, done = env.step(action)

        agent.buffer.push(state, action, reward, next_state, done)
        agent.train()

        state = next_state
        total_reward += reward

    rewards_history.append(total_reward)

    # Update the target network
    if (episode + 1) % 10 == 0:
        agent.update_target_network()

    if (episode + 1) % 50 == 0:
        avg_reward = np.mean(rewards_history[-50:])
        print(f"Episode {episode+1}: Avg Reward = {avg_reward:.2f}, epsilon = {agent.epsilon:.3f}")

# Learning curve
plt.figure(figsize=(10, 6))
plt.plot(np.convolve(rewards_history, np.ones(20)/20, mode='valid'))
plt.xlabel('Episode')
plt.ylabel('Average Reward (20 episodes)')
plt.title('DQN: Learning Progress in Continuous-State Materials Discovery')
plt.grid(True)
plt.show()

Example output:

Episode 50: Avg Reward = -45.23, epsilon = 0.779
Episode 100: Avg Reward = -32.15, epsilon = 0.606
Episode 200: Avg Reward = -18.92, epsilon = 0.365
Episode 500: Avg Reward = -8.45, epsilon = 0.010

Explanation: - The neural network learns the Q-function for continuous states - epsilon decays, shifting from exploration to exploitation - Eventually, materials close to the target properties are found efficiently

Exercises

Problem 1 (Difficulty: easy)

In the Q-learning update rule, explain what happens when you increase the learning rate $\alpha$. Also, describe what happens in the extreme cases of $\alpha=0$ and $\alpha=1$.

Problem 2 (Difficulty: medium)

In materials discovery, a reward function was designed as follows. State the problems with this design and propose improvements.

def reward_function(material_property, target=3.0):
    if material_property == target:
        return 1.0
    else:
        return 0.0

def improved_reward_function(material_property, target=3.0):
    # Continuous reward based on distance from the target
    distance = abs(material_property - target)

    if distance < 0.1:
        return 10.0  # Very close (bonus)
    elif distance < 0.5:
        return 5.0   # Close
    else:
        return -distance  # The farther, the greater the penalty

Problem 3 (Difficulty: hard)

Explain the roles of "experience replay" and the "target network" in DQN, and use Python code to experiment with what problems occur when each is absent.

# Version without experience replay
class DQNAgentNoReplay(DQNAgent):
    def train_no_replay(self, state, action, reward, next_state, done):
        # Train on only the latest transition
        states = torch.FloatTensor([state])
        actions = torch.LongTensor([action]).unsqueeze(1)
        rewards = torch.FloatTensor([reward]).unsqueeze(1)
        next_states = torch.FloatTensor([next_state])
        dones = torch.FloatTensor([done]).unsqueeze(1)

        current_q = self.policy_net(states).gather(1, actions)
        with torch.no_grad():
            max_next_q = self.target_net(next_states).max(1)[0].unsqueeze(1)
            target_q = rewards + (1 - dones) * self.gamma * max_next_q

        loss = nn.MSELoss()(current_q, target_q)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

# Version without a target network (uses policy_net for the TD target)
# -> Learning becomes unstable

# Result: Without experience replay, convergence is slow; without the target network, it oscillates

Summary of This Section

• The materials search space is vast, and conventional trial and error is inefficient
• Reinforcement learning learns an optimal search strategy through interaction with an environment
• The Markov Decision Process (MDP) is the mathematical foundation of reinforcement learning
• Q-learning is effective for discrete states and actions, recording values in a Q-table
• DQN approximates the Q-function with a neural network, handling huge state spaces
• Experience replay and the target network stabilize DQN learning

In the next chapter, we will learn more advanced policy gradient methods and Actor-Critic techniques.

Quality Checklist: Verifying Your Materials Discovery RL Implementation

MDP Formulation Skills

• [ ] Can formulate a materials discovery task in terms of states, actions, and rewards
• [ ] Can judge whether the Markov property assumption is valid
• [ ] Can verify that the reward function correctly represents the search objective
• [ ] Can explain the rationale for choosing the discount factor gamma (usually 0.95-0.99 in materials discovery)

Q-Learning Implementation Skills

• [ ] Can implement epsilon-greedy exploration
• [ ] Can implement TD error computation and Q-value updates
• [ ] Can judge convergence from the learning curve
• [ ] Understand the size constraints of the Q-table (number of states x actions)

DQN Implementation Skills

• [ ] Can define a neural network in PyTorch
• [ ] Can implement an experience replay buffer
• [ ] Can set the timing for updating the target network
• [ ] Can appropriately set the epsilon decay schedule

Considerations Specific to Materials Discovery

• [ ] Understand the difference between composition-based and structure-based descriptors
• [ ] Can design a state that distinguishes material polymorphs
• [ ] Can incorporate physical constraints of material properties into the reward function
• [ ] Can design a search strategy that accounts for DFT computation costs

Debugging Skills

• [ ] Can identify the cause when Q-values diverge
• [ ] Know how to handle cases where exploration stalls
• [ ] Can detect and fix reward scaling problems
• [ ] Can systematically investigate the impact of hyperparameters

Code Quality

• [ ] Document dependency library versions in all code
• [ ] Fix the random seed to ensure reproducibility
• [ ] Write validation code for data shape, type, and range
• [ ] Implement error handling (exception handling)

Next Steps

If your achievement is below 80%: - Re-read this chapter and redo the exercises - Get hands-on implementation experience with a simple grid world

If your achievement is 80-95%: - You are ready to proceed to Chapter 2 (policy gradient methods) - Try implementing the DQN code from scratch yourself

If your achievement is 95% or higher: - Proceed to Chapter 2 and learn more advanced methods - Try RL on an actual materials discovery task

References

1. Mnih et al. "Playing Atari with Deep Reinforcement Learning" arXiv (2013) - The original DQN paper
2. Sutton & Barto "Reinforcement Learning: An Introduction" MIT Press (2018) - RL textbook
3. Zhou et al. "Optimization of molecules via deep reinforcement learning" Scientific Reports (2019)
4. Ling et al. "High-dimensional materials and process optimization using data-driven experimental design with well-calibrated uncertainty estimates" Integrating Materials and Manufacturing Innovation (2017)

Next Chapter: Chapter 2: Theoretical Foundations of Reinforcement Learning