Chapter 3: Building Materials Discovery Environments

We illustrate the definitions of states, actions, and rewards with examples, and make concrete how to construct an exploration strategy. We also review key points for leveraging simulators.

💡 Supplement: Estimate the simulator's "gap from reality" and make it robust with Domain Randomization.

Learning Objectives

In this chapter, you will master the following:

• How to implement custom OpenAI Gym environments
• Design of material descriptors and state spaces
• Design principles for effective reward functions
• How to integrate with DFT calculations and experimental instruments

3.1 Fundamentals of OpenAI Gym Environments

Components of a Gym Environment

OpenAI Gym is the standard interface for reinforcement learning environments. Every Gym environment implements the following methods:

import gym
import numpy as np

class CustomEnv(gym.Env):
    """Template for a custom Gym environment"""

    def __init__(self):
        super(CustomEnv, self).__init__()

        # Define the action space and observation space (required)
        self.action_space = gym.spaces.Discrete(4)  # Discrete actions (4 types)
        self.observation_space = gym.spaces.Box(
            low=0, high=10, shape=(4,), dtype=np.float32
        )  # Continuous state (4 dimensions, range [0, 10])

    def reset(self):
        """Reset the environment to its initial state

        Returns:
            observation: initial state
        """
        self.state = np.random.uniform(0, 10, 4).astype(np.float32)
        return self.state

    def step(self, action):
        """Execute an action and advance the environment by one step

        Args:
            action: the action to execute

        Returns:
            observation: next state
            reward: reward
            done: episode termination flag
            info: additional information (dictionary)
        """
        # Update the state according to the action
        self.state = self._update_state(action)

        # Compute the reward
        reward = self._compute_reward()

        # Check the termination condition
        done = self._is_done()

        # Additional information
        info = {'distance': self._compute_distance()}

        return self.state, reward, done, info

    def render(self, mode='human'):
        """Visualize the environment (optional)"""
        print(f"Current state: {self.state}")

    def _update_state(self, action):
        """State update logic"""
        # Implementation depends on the environment
        pass

    def _compute_reward(self):
        """Reward computation logic"""
        pass

    def _is_done(self):
        """Termination condition check"""
        pass

    def _compute_distance(self):
        """Computation of additional information"""
        pass

Defining Action and Observation Spaces

Gym supports a variety of space types:

from gym import spaces

# Discrete actions (integers 0, 1, 2, 3)
action_space = spaces.Discrete(4)

# Continuous actions (real-valued vector [-1, 1]^3)
action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)

# Dictionary form (multiple inputs)
observation_space = spaces.Dict({
    'composition': spaces.Box(low=0, high=1, shape=(10,), dtype=np.float32),
    'temperature': spaces.Box(low=0, high=1000, shape=(1,), dtype=np.float32),
    'pressure': spaces.Box(low=0, high=10, shape=(1,), dtype=np.float32)
})

# Tuple form
action_space = spaces.Tuple((
    spaces.Discrete(5),      # Element selection
    spaces.Box(low=0, high=1, shape=(1,))  # Composition ratio
))

# Multi-binary (multiple binary selections)
action_space = spaces.MultiBinary(10)  # Turn 10 elements ON/OFF

3.2 Design of Material Descriptors and State Spaces

Selecting Material Descriptors

The state space is a representation of a material's properties as a numerical vector. Choosing effective descriptors is important.

1. Composition-Based Descriptors

Element fractions:

# Example: composition vector of Li2MnO3
composition = {
    'Li': 2/6,   # 33.3%
    'Mn': 1/6,   # 16.7%
    'O': 3/6     # 50.0%
}

# Vector over the entire periodic table (118 dimensions)
state = np.zeros(118)
state[2] = 0.333   # Li (atomic number 3)
state[24] = 0.167  # Mn (atomic number 25)
state[7] = 0.500   # O (atomic number 8)

Magpie descriptors (Ward et al., 2016):

from matminer.featurizers.composition import ElementProperty

featurizer = ElementProperty.from_preset("magpie")
# Generate a 132-dimensional descriptor from composition
# - mean atomic number, mean electronegativity, mean ionic radius, etc.
composition = "Li2MnO3"
features = featurizer.featurize(Composition(composition))

2. Structure-Based Descriptors

Lattice constants:

# Crystal lattice
state = np.array([
    a, b, c,           # Lattice constants
    alpha, beta, gamma # Angles
])

Smooth Overlap of Atomic Positions (SOAP):

from dscribe.descriptors import SOAP
from ase import Atoms

# Generate descriptor from atomic structure
atoms = Atoms('H2O', positions=[[0, 0, 0], [0, 0, 1], [0, 1, 0]])
soap = SOAP(species=['H', 'O'], rcut=5.0, nmax=8, lmax=6)
state = soap.create(atoms)  # High-dimensional vector

3. Process Parameters

Synthesis conditions:

# State of the synthesis process
state = np.array([
    temperature,      # Temperature [K]
    pressure,         # Pressure [Pa]
    time,             # Time [s]
    heating_rate,     # Heating rate [K/min]
    atmosphere_O2     # Oxygen partial pressure [Pa]
])

Worked Example: Bandgap Discovery Environment

from pymatgen.core import Composition
from matminer.featurizers.composition import ElementProperty

class BandgapDiscoveryEnv(gym.Env):
    """Bandgap optimization environment

    Goal: discover a material with a specific bandgap (e.g., 3.0 eV)
    """

    def __init__(self, target_bandgap=3.0, element_pool=None):
        super(BandgapDiscoveryEnv, self).__init__()

        self.target_bandgap = target_bandgap

        # Available elements (default: typical semiconductor elements)
        if element_pool is None:
            self.element_pool = ['Ti', 'Zr', 'Hf', 'V', 'Nb', 'Ta', 'Cr', 'Mo', 'W',
                                  'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge',
                                  'As', 'Se', 'Sr', 'Y', 'In', 'Sn', 'Sb', 'Te', 'O', 'S', 'N']
        else:
            self.element_pool = element_pool

        self.n_elements = len(self.element_pool)

        # Action space: select 3 elements + ratio of each element
        # Simplified: discrete selection of 3 elements (combination)
        self.action_space = gym.spaces.MultiDiscrete([self.n_elements] * 3)

        # State space: Magpie descriptor (132 dimensions)
        self.featurizer = ElementProperty.from_preset("magpie")
        self.observation_space = gym.spaces.Box(
            low=-10, high=10, shape=(132,), dtype=np.float32
        )

        # History (compositions tried)
        self.history = []
        self.current_composition = None

    def reset(self):
        """Random initial composition"""
        self.history = []
        action = self.action_space.sample()
        self.current_composition = self._action_to_composition(action)
        return self._get_state()

    def step(self, action):
        """Try a new material composition"""
        self.current_composition = self._action_to_composition(action)

        # State (descriptor)
        state = self._get_state()

        # Predict the bandgap (surrogate model or DFT)
        predicted_bandgap = self._predict_bandgap(self.current_composition)

        # Reward: negative of the difference from the target
        error = abs(predicted_bandgap - self.target_bandgap)
        reward = -error

        # Bonus reward (when close to the target)
        if error < 0.1:
            reward += 10.0  # Very close

        # Add to history
        self.history.append({
            'composition': self.current_composition,
            'bandgap': predicted_bandgap,
            'reward': reward
        })

        # Termination condition: reached the target or max steps
        done = error < 0.05 or len(self.history) >= 100

        info = {
            'composition': self.current_composition,
            'predicted_bandgap': predicted_bandgap,
            'error': error
        }

        return state, reward, done, info

    def _action_to_composition(self, action):
        """Convert an action into a composition string

        Args:
            action: [elem1_idx, elem2_idx, elem3_idx]

        Returns:
            composition string (e.g., "TiO2")
        """
        elements = [self.element_pool[idx] for idx in action]

        # Remove duplicates
        unique_elements = list(set(elements))

        # Simplified: equal-amount mixing
        if len(unique_elements) == 1:
            comp_str = unique_elements[0]
        elif len(unique_elements) == 2:
            comp_str = f"{unique_elements[0]}{unique_elements[1]}"
        else:
            comp_str = f"{unique_elements[0]}{unique_elements[1]}{unique_elements[2]}"

        return comp_str

    def _get_state(self):
        """Generate a descriptor from the current composition"""
        try:
            comp = Composition(self.current_composition)
            features = self.featurizer.featurize(comp)
            return np.array(features, dtype=np.float32)
        except:
            # For an invalid composition, return a zero vector
            return np.zeros(132, dtype=np.float32)

    def _predict_bandgap(self, composition):
        """Predict the bandgap

        In practice:
        - a machine learning model (pre-trained)
        - a DFT calculation (pymatgen + VASP)
        - a database lookup (Materials Project)

        Here we use a simple rule-based approach
        """
        try:
            comp = Composition(composition)

            # Simple rule: oxygen-containing compounds tend to have larger bandgaps
            if 'O' in comp:
                base_gap = 2.5
            elif 'S' in comp:
                base_gap = 1.8
            elif 'N' in comp:
                base_gap = 2.0
            else:
                base_gap = 1.0

            # Influence of metallic elements
            metals = ['Ti', 'Zr', 'Hf', 'V', 'Nb', 'Ta']
            for metal in metals:
                if metal in comp:
                    base_gap += 0.5

            # Random noise (experimental error)
            noise = np.random.normal(0, 0.2)
            return max(0, base_gap + noise)

        except:
            return 0.0

    def render(self, mode='human'):
        print(f"Current composition: {self.current_composition}")
        if self.history:
            last = self.history[-1]
            print(f"Predicted bandgap: {last['bandgap']:.2f} eV")
            print(f"Target: {self.target_bandgap:.2f} eV")
            print(f"Reward: {last['reward']:.2f}")


# Test the environment
env = BandgapDiscoveryEnv(target_bandgap=3.0)

state = env.reset()
print(f"Initial state: {state.shape}")

for step in range(10):
    action = env.action_space.sample()
    state, reward, done, info = env.step(action)

    print(f"\nStep {step+1}:")
    print(f"  Composition: {info['composition']}")
    print(f"  Predicted bandgap: {info['predicted_bandgap']:.2f} eV")
    print(f"  Reward: {reward:.2f}")

    if done:
        print("Target reached!")
        break

Example output:

Initial state: (132,)

Step 1:
  Composition: TiO
  Predicted bandgap: 3.12 eV
  Reward: -0.12

Step 2:
  Composition: ZrO
  Predicted bandgap: 2.95 eV
  Reward: -0.05
Target reached!

3.3 Designing Effective Reward Functions

Principles of Reward Design

The reward function defines what the agent should optimize. Inappropriate rewards cause undesirable behavior or learning failure.

Principle 1: A Clear Goal

Bad example:

# Ambiguous reward
reward = 1 if 'good_material' else 0  # the definition of "good" is unclear

Good example:

# Clear goal (bandgap)
target = 3.0
predicted = 2.8
reward = -abs(predicted - target)  # distance from the target

Principle 2: Scaling

Set the range of the reward appropriately:

Bad example:

# The reward is extremely large
reward = 1e10 if success else -1e10  # learning becomes unstable

Good example:

# Normalize to roughly [-1, 1]
reward = -error / max_error  # error ∈ [0, max_error]

Principle 3: Shaping (Intermediate Rewards)

Convert sparse rewards into dense rewards:

Sparse reward (hard to learn):

reward = 1.0 if distance < 0.1 else 0.0

Dense reward (easy to learn):

# A continuous reward that depends on distance
reward = -distance

# A further hierarchical reward
if distance < 0.5:
    reward += 5.0  # close
if distance < 0.1:
    reward += 10.0  # very close

Principle 4: Multi-Objective Optimization

Weight multiple objectives:

# Optimize both bandgap and stability
bandgap_error = abs(predicted_bandgap - target_bandgap)
stability = formation_energy  # a negative value is stable

# Weighted reward
w1, w2 = 0.7, 0.3
reward = -w1 * bandgap_error - w2 * max(0, stability)

Worked Examples of Reward Design

Example 1: Maximizing Catalytic Activity

class CatalystOptimizationEnv(gym.Env):
    """Environment for maximizing catalytic activity"""

    def _compute_reward(self, activity, selectivity, stability):
        """Multi-objective reward

        Args:
            activity: catalytic activity (higher is better)
            selectivity: selectivity (selectivity toward the target product, higher is better)
            stability: stability (negative formation energy, lower is more stable)

        Returns:
            overall reward
        """
        # Normalize each metric to [0, 1]
        activity_norm = activity / 100.0  # assume max of 100
        selectivity_norm = selectivity  # already in [0, 1]
        stability_norm = -stability / 5.0  # assume max of -5 eV

        # Weighted sum (emphasizing activity)
        weights = {'activity': 0.5, 'selectivity': 0.3, 'stability': 0.2}
        reward = (weights['activity'] * activity_norm +
                  weights['selectivity'] * selectivity_norm +
                  weights['stability'] * stability_norm)

        # Penalty: unstable material
        if stability > 0:  # positive formation energy (unstable)
            reward -= 1.0

        return reward

Example 2: Synthesis Cost Constraint

def reward_with_cost_constraint(self, performance, synthesis_cost, max_cost=1000):
    """Reward with a cost constraint

    Args:
        performance: material performance
        synthesis_cost: synthesis cost [USD/kg]
        max_cost: cost upper limit

    Returns:
        reward
    """
    # Base reward based on performance
    base_reward = performance

    # Penalty for violating the cost constraint
    if synthesis_cost > max_cost:
        penalty = (synthesis_cost - max_cost) / max_cost
        base_reward -= 10.0 * penalty

    # Bonus for lower cost
    cost_bonus = max(0, (max_cost - synthesis_cost) / max_cost)
    base_reward += 2.0 * cost_bonus

    return base_reward

3.4 Integration with DFT Calculations

Retrieving Data from the Materials Project

Retrieve real material properties and use them in the reward:

from mp_api.client import MPRester
import os

class MPIntegratedEnv(gym.Env):
    """Materials Project integrated environment"""

    def __init__(self, mp_api_key=None):
        super(MPIntegratedEnv, self).__init__()

        # Materials Project API key
        if mp_api_key is None:
            mp_api_key = os.getenv("MP_API_KEY")

        self.mpr = MPRester(mp_api_key)

        # ... (environment setup) ...

    def _get_bandgap_from_mp(self, composition):
        """Retrieve the bandgap from the Materials Project

        Args:
            composition: composition (e.g., "TiO2")

        Returns:
            bandgap [eV] (None if no data is available)
        """
        try:
            # Search by composition
            docs = self.mpr.materials.summary.search(
                formula=composition,
                fields=["material_id", "band_gap", "formation_energy_per_atom"]
            )

            if docs:
                # Select the most stable structure (minimum formation energy)
                stable_doc = min(docs, key=lambda x: x.formation_energy_per_atom)
                return stable_doc.band_gap
            else:
                return None

        except Exception as e:
            print(f"Materials Project search error: {e}")
            return None

    def step(self, action):
        composition = self._action_to_composition(action)

        # Retrieve data from the Materials Project
        bandgap = self._get_bandgap_from_mp(composition)

        if bandgap is not None:
            # Compute the reward with real data
            error = abs(bandgap - self.target_bandgap)
            reward = -error
        else:
            # If no data, use a prediction model or apply a penalty
            reward = -10.0  # penalty for an unknown material

        # ... (state, termination condition, etc.) ...

        return state, reward, done, info

Note: Avoid sending a large number of requests to the Materials Project, and make use of a local cache.

Integrating DFT Calculations with ASE (Advanced)

from ase import Atoms
from ase.calculators.vasp import Vasp
from ase.optimize import BFGS

class DFTIntegratedEnv(gym.Env):
    """DFT-integrated environment (high computational cost)"""

    def _calculate_bandgap_dft(self, composition):
        """Obtain the bandgap via a DFT calculation

        Warning: this is very time-consuming (hours to days per material)
        In practice, use a precomputed database

        Args:
            composition: composition

        Returns:
            bandgap [eV]
        """
        # Generate the crystal structure (with pymatgen, etc.)
        structure = self._generate_structure(composition)

        # Convert to an ASE Atoms object
        atoms = Atoms(
            symbols=structure.species,
            positions=structure.cart_coords,
            cell=structure.lattice.matrix,
            pbc=True
        )

        # VASP calculation settings
        calc = Vasp(
            xc='PBE',
            encut=520,
            kpts=(4, 4, 4),
            ismear=0,
            sigma=0.05,
            directory='vasp_calc'
        )
        atoms.calc = calc

        # Structural optimization
        opt = BFGS(atoms)
        opt.run(fmax=0.05)

        # Bandgap calculation
        # ... (parsing the VASP OUTCAR) ...

        return bandgap

    def step(self, action):
        # Because DFT calculations are time-consuming,
        # in practice the following measures are needed:
        # 1. Build a precomputed database
        # 2. Predict quickly with a surrogate model
        # 3. Run DFT only for important materials via active learning
        pass

Practical approach: 1. Pre-training: train a surrogate model on data such as the Materials Project 2. Reinforcement learning: explore quickly with the surrogate model 3. Validation: precisely evaluate only promising materials with DFT calculations

3.5 Integration with Experimental Instruments (Closed Loop)

Controlling Automated Experimental Instruments via a REST API

import requests

class RoboticLabEnv(gym.Env):
    """Robotic experimental instrument integrated environment"""

    def __init__(self, api_endpoint="http://lab-robot.example.com/api"):
        super(RoboticLabEnv, self).__init__()
        self.api_endpoint = api_endpoint

        # ... (environment setup) ...

    def _synthesize_and_measure(self, composition, temperature, time):
        """Synthesize a material and measure its properties

        Args:
            composition: composition
            temperature: synthesis temperature [K]
            time: synthesis time [min]

        Returns:
            measurement results (bandgap, XRD pattern, etc.)
        """
        # Send a synthesis request to the robot
        payload = {
            'composition': composition,
            'temperature': temperature,
            'time': time,
            'measurement': ['bandgap', 'xrd']
        }

        response = requests.post(
            f"{self.api_endpoint}/synthesize",
            json=payload,
            headers={'Authorization': 'Bearer YOUR_API_KEY'}
        )

        if response.status_code == 200:
            result = response.json()
            return result['bandgap'], result['xrd_pattern']
        else:
            raise Exception(f"Experiment failed: {response.text}")

    def step(self, action):
        """Action = synthesis conditions"""
        composition, temperature, time = self._decode_action(action)

        # Run the experiment (minutes to hours)
        bandgap, xrd = self._synthesize_and_measure(composition, temperature, time)

        # Compute the reward
        reward = -abs(bandgap - self.target_bandgap)

        # Update the state (including the experimental history)
        state = self._update_state(composition, temperature, time, bandgap, xrd)

        done = len(self.history) >= self.max_experiments

        return state, reward, done, {'bandgap': bandgap}

Challenges: - Experimental cost: thousands to tens of thousands of yen per run - Time: hours to days for synthesis and measurement - Safety: robot malfunctions, handling of hazardous substances

Solutions: - Simulation first: explore in advance with a surrogate model - Combine with Bayesian optimization: efficiently select experimental points - Batch experiments: synthesize multiple materials in parallel

Exercises

Problem 1 (Difficulty: easy)

Explain the difference between the following two reward functions, and state which is easier to learn along with the reason.

Reward A:

reward = 10.0 if abs(bandgap - 3.0) < 0.1 else 0.0

Reward B:

reward = -abs(bandgap - 3.0)

# Hybrid reward
reward = -abs(bandgap - 3.0)  # dense reward
if abs(bandgap - 3.0) < 0.1:
    reward += 10.0  # bonus (the sparse-reward element)

Problem 2 (Difficulty: medium)

In materials discovery, compare the following three state representations and describe the advantages and disadvantages of each.

1. Composition only: ["Li2MnO3"] (a string)
2. Element fractions: [0.33, 0.17, 0.50] (fractions of Li, Mn, O)
3. Magpie descriptor: a 132-dimensional vector (mean atomic number, electronegativity, etc.)

Problem 3 (Difficulty: hard)

In the bandgap discovery environment, implement the following improvements:

1. History-aware state: include information about materials tried so far in the state
2. Exploration bonus: an additional reward when exploring an unknown region
3. Early termination: end the episode if there is no improvement for 10 consecutive steps

import numpy as np
from scipy.spatial.distance import euclidean

class ImprovedBandgapEnv(gym.Env):
    """Improved bandgap discovery environment"""

    def __init__(self, target_bandgap=3.0):
        super(ImprovedBandgapEnv, self).__init__()

        self.target_bandgap = target_bandgap

        # Action/state spaces (simplified)
        self.action_space = gym.spaces.Box(low=0, high=1, shape=(10,), dtype=np.float32)
        self.observation_space = gym.spaces.Box(low=-10, high=10, shape=(15,), dtype=np.float32)

        # History
        self.history = []
        self.best_error = float('inf')
        self.no_improvement_count = 0

    def reset(self):
        self.history = []
        self.best_error = float('inf')
        self.no_improvement_count = 0

        initial_state = self._get_state(np.random.uniform(0, 1, 10))
        return initial_state

    def step(self, action):
        # Bandgap prediction (simple model)
        predicted_bandgap = np.sum(action) * 3.0  # placeholder prediction

        # Error
        error = abs(predicted_bandgap - self.target_bandgap)

        # Base reward
        reward = -error

        # Improvement 1: history-aware state
        state = self._get_state(action)

        # Improvement 2: exploration bonus
        exploration_bonus = self._compute_exploration_bonus(action)
        reward += 0.1 * exploration_bonus

        # Improvement 3: early termination
        if error < self.best_error:
            self.best_error = error
            self.no_improvement_count = 0
        else:
            self.no_improvement_count += 1

        done = error < 0.05 or self.no_improvement_count >= 10 or len(self.history) >= 100

        # Add to history
        self.history.append({
            'action': action,
            'bandgap': predicted_bandgap,
            'error': error
        })

        info = {'bandgap': predicted_bandgap, 'exploration_bonus': exploration_bonus}

        return state, reward, done, info

    def _get_state(self, action):
        """History-aware state

        State composition:
        - current action (10 dimensions)
        - distance to the best material (1 dimension)
        - history size (1 dimension)
        - number of consecutive non-improvements (1 dimension)
        - mean error (1 dimension)
        - best error (1 dimension)
        """
        state = np.zeros(15, dtype=np.float32)

        # Current action
        state[:10] = action

        # Distance to the best material
        if self.history:
            best_action = min(self.history, key=lambda x: x['error'])['action']
            state[10] = euclidean(action, best_action) / 10.0  # normalize
        else:
            state[10] = 1.0

        # History size
        state[11] = len(self.history) / 100.0  # normalize

        # Number of consecutive non-improvements
        state[12] = self.no_improvement_count / 10.0

        # Mean error
        if self.history:
            state[13] = np.mean([h['error'] for h in self.history])
        else:
            state[13] = 10.0

        # Best error
        state[14] = self.best_error

        return state

    def _compute_exploration_bonus(self, action):
        """Exploration bonus

        The farther from past actions, the higher the bonus
        """
        if not self.history:
            return 1.0  # always explore at the start

        # Minimum distance to past actions
        min_distance = min(
            euclidean(action, h['action'])
            for h in self.history
        )

        # The larger the distance, the higher the bonus (max 1.0)
        bonus = min(1.0, min_distance / 5.0)

        return bonus


# Test
env = ImprovedBandgapEnv()
state = env.reset()

for step in range(50):
    action = env.action_space.sample()
    state, reward, done, info = env.step(action)

    print(f"Step {step+1}: Bandgap={info['bandgap']:.2f}, "
          f"Reward={reward:.2f}, Exploration={info['exploration_bonus']:.2f}")

    if done:
        print(f"Terminated: best error={env.best_error:.4f}, "
              f"consecutive non-improvements={env.no_improvement_count}")
        break

Summary of This Section

• OpenAI Gym is the standard interface for reinforcement learning environments
• The state space is designed with material descriptors (composition, structure, process parameters)
• The reward function needs a clear goal, appropriate scaling, and intermediate rewards
• DFT integration is accelerated with a surrogate model, running precise calculations only for important materials
• Experimental instrument integration achieves closed-loop optimization via a REST API

In the next chapter, you will learn about real-world application cases such as chemical process control and synthesis route design.

References

1. Brockman et al. "OpenAI Gym" arXiv (2016) - the standard for Gym environments
2. Ward et al. "A general-purpose machine learning framework for predicting properties of inorganic materials" npj Computational Materials (2016) - Magpie descriptors
3. Brockherde et al. "Bypassing the Kohn-Sham equations with machine learning" Nature Communications (2017) - DFT acceleration
4. Ng et al. "Policy invariance under reward transformations" ICML (1999) - reward shaping theory

Next chapter: Chapter 4: Real-World Applications and Closed Loops