Application of Reinforcement Learning for Multigroup Energy Grid Optimization for Neutron Transport Criticality Problems

This paper presents a reinforcement learning approach combined with neural network surrogate models to optimize multigroup energy structures for one-dimensional spherical k-criticality neutron transport problems, achieving accuracy comparable to or better than existing methods while offering greater flexibility and computational efficiency.

The Gist

The Big Picture: Tuning the Radio for a Nuclear Signal

Imagine you are trying to listen to a very faint radio signal coming from a nuclear reactor. The signal (neutrons) is complex, with different "frequencies" (energies) that change rapidly. To understand the signal, you need to tune your radio dial.

In nuclear physics, scientists use a method called Multigroup Neutron Transport. Think of this as dividing the entire radio spectrum into a set number of "channels" or "bins" (called energy groups).

• Too many bins: You get a crystal-clear picture of the signal, but your computer has to do so much work that it takes days to finish the calculation. It's like trying to listen to every single frequency individually.
• Too few bins: The computer runs fast, but you might miss important details or hear static, leading to inaccurate results.

The goal of this paper is to find the perfect number of bins and the perfect places to draw the lines between them for a specific nuclear problem.

The Problem: The "Goldilocks" Dilemma

For decades, scientists have used standard "preset" channel layouts (like the LANL30 or LANL70 structures). These are like buying a radio with fixed buttons. They work okay for many situations, but they aren't perfect for every specific reactor.

Finding the best custom layout is hard.

1. It's expensive: To test if a new layout works, you have to run a massive, slow computer simulation (like running a full physics test for every single button press).
2. It's tricky: If you just start guessing, you might get stuck in a "local minimum." Imagine you are in a foggy valley; you might think you've reached the bottom because you can't see the deeper valley just over the next hill.

The Solution: A Smart Robot with a Crystal Ball

The authors, Ben Whewell and his team at Los Alamos National Laboratory, used Reinforcement Learning (RL).

The Analogy:
Imagine a robot trying to solve a maze.

• The Robot (RL Agent): Its job is to start with a very detailed map (a high-fidelity grid with 618 channels) and remove lines until it reaches a target number (like 30 or 70).
• The Reward: Every time the robot removes a line, it gets a score. It wants a high score, which means the simulation is still accurate and it has removed as many lines as possible to save time.
• The Trap: If the robot just guesses, it will take millions of tries to learn, and each try requires a slow, expensive physics simulation.

The Secret Weapon: The Surrogate Model (The Crystal Ball)
To make the robot learn faster, the team built a Neural Network Surrogate Model.

• Think of this as a crystal ball or a highly experienced coach.
• Instead of running the slow, expensive physics simulation every time the robot makes a move, the robot asks the crystal ball: "If I remove this line, how good will the result be?"
• The crystal ball looks at the pattern of the lines and the materials (like Uranium or Plutonium) and instantly predicts the accuracy. It doesn't give a perfect number, but it puts the result into a "quality bucket" (e.g., "This is a 9 out of 10").

This allows the robot to practice millions of times in a few hours instead of thousands of years.

What They Did

They tested this "Robot + Crystal Ball" team on two famous nuclear puzzles:

1. Godiva: A sphere of pure Uranium.
2. BeRP Ball: A sphere of Plutonium surrounded by a shell of Beryllium.

They taught the robot to start with a massive grid and "prune" it down to 30 or 70 groups, learning which lines were essential to keep and which could be cut.

The Results: Better than the Standard

When they tested the robot's custom layouts against the standard "preset" layouts (LANL30 and LANL70):

• Accuracy: The robot's custom layouts were more accurate. They captured the important details of the nuclear reaction better than the standard presets.
• Speed: The robot learned to find these good layouts much faster than previous methods (like "Hierarchical Agglomeration," which is a slow, step-by-step greedy approach).
• Flexibility: The robot learned a general strategy. If you changed the size of the sphere or the material, the robot could adapt without needing to be retrained from scratch.

Key Takeaways in Plain English

1. Smart Pruning: Instead of building a grid from scratch, the AI starts with a perfect, detailed grid and learns exactly which parts to cut away to save time without losing accuracy.
2. The Coach: They used a fast AI "coach" (surrogate model) to predict results, saving them from running slow, expensive simulations millions of times.
3. Winning: The AI-designed grids beat the old, standard grids for these specific nuclear tests, offering a more flexible and efficient way to solve nuclear physics problems.

In short, they taught a computer to be a master tuner, finding the perfect balance between speed and accuracy for nuclear safety calculations, using a "crystal ball" to speed up the learning process.

Technical Summary

Technical Summary: Application of Reinforcement Learning for Multigroup Energy Grid Optimization for Neutron Transport Criticality Problems

Problem Statement
Accurate neutron transport calculations rely heavily on the multigroup discretization scheme, where the continuous energy variable is integrated over finite ranges to create piecewise constant energy groups. The selection of energy group bounds is critical; suboptimal bounds can lead to significant errors in neutron flux spectra and reaction rates. While high-fidelity grids (e.g., LANL618) offer accuracy, they incur high computational costs and memory footprints. Conversely, low-fidelity grids (e.g., LANL30, LANL70) reduce costs but require careful selection of bounds to maintain accuracy. Existing optimization techniques, such as Particle Swarm Optimization (PSO) and Hierarchical Agglomeration (HA), face challenges including high computational costs due to the requirement of full transport simulations for every evaluation step and susceptibility to local minima or poor convergence.

Methodology
The authors propose a novel framework combining Reinforcement Learning (RL) with neural network surrogate modeling to optimize energy group structures for one-dimensional spherical $k$-criticality problems.

• Reinforcement Learning Formulation: The problem is modeled using the Proximal Policy Optimization (PPO) algorithm.

  • State Space: A binary vector of length 619 representing the presence or absence of energy bounds from a reference LANL618 grid. For non-homogeneous problems (e.g., BeRP ball), material thickness and total cross-section data are appended.
  • Action Space: The agent removes one energy bound at a time, transitioning from a high-fidelity starting state ($G_{max} \in [200, 617]$) toward a target number of groups ($G_{min}$). Action masking ensures only valid removals occur.
  • Reward Function: The reward balances two objectives: minimizing the number of energy groups and maximizing grid accuracy. Accuracy is evaluated via an error metric ($\epsilon$) that combines the relative errors of the effective multiplication factor ($k_{eff}$) and integrated reaction rates (total, $\nu$-fission, and absorption). To prevent error cancellation from masking flux inaccuracies, the $k_{eff}$ error is weighted by a factor of 3 in the root-sum-square calculation.

• Surrogate Modeling: To overcome the sample inefficiency of on-policy RL (which would otherwise require millions of full transport simulations), a 10-class classification neural network surrogate model is employed.

  • Architecture: For homogeneous problems (Godiva), a 1D Convolutional Neural Network (CNN) processes the binary energy bound vector. For heterogeneous problems (BeRP ball), a multimodal architecture combines the CNN with a Long Short-Term Memory (LSTM) network to encode spatial and material properties.
  • Training Data: Random subsets of the LANL618 grid are generated, and full transport simulations are run to calculate the error metric $\epsilon$. These errors are transformed into normal distributions and binned into 10 classes (1 = least accurate, 10 = most accurate).
  • Integration: The surrogate model outputs the probability distribution over these 10 classes. The expected class value is used to calculate the reward, allowing the RL agent to learn without executing a full transport simulation at every step.

Key Contributions

1. RL for Group Structure Optimization: This work introduces the application of PPO-based RL to the specific problem of optimizing multigroup energy structures, allowing the agent to identify critical bounds without being restricted to a fixed initial grid topology (beyond the LANL618 subset constraint).
2. Surrogate-Accelerated Training: The development of a classification-based surrogate model that incorporates energy, material, and spatial information significantly reduces the computational cost of RL training, replacing expensive transport simulations with fast neural network inference.
3. Flexible Optimization: Unlike greedy hierarchical methods that require re-running simulations for every new starting condition, the trained RL agents can adapt to different starting group structures and material layouts without retraining.

Results
The method was validated on two benchmark problems: the Godiva (uranium sphere) and the BeRP ball (plutonium sphere with beryllium reflector).

• Surrogate Performance:
  • Godiva: The CNN surrogate achieved 78.3% true accuracy and 98.2% adjacent accuracy (prediction within one class) on subcritical test data, generalizing well to supercritical configurations.
  • BeRP Ball: The multimodal CNN-LSTM surrogate achieved 70.8% true accuracy and 97.4% adjacent accuracy across varying plutonium radii and criticality states.
• RL Optimization Performance:
  • Accuracy: The RL-constructed group structures (RL30 and RL70) outperformed standard LANL30 and LANL70 structures in terms of both $k_{eff}$ and reaction rate errors when compared against the LANL618 reference.
  • Comparison with HA: The RL method achieved performance comparable to the Hierarchical Agglomeration (HA) method but with significantly reduced computational overhead. While HA required tens of thousands of full transport simulations (45,225 for 301-group start, 191,362 for LANL618 start) to optimize a single problem, the RL method required only two trained models (for targets of 30 and 70 groups) and no retraining for different starting grids or material layouts.
  • Training Efficiency: Training the surrogate and RL models took approximately six hours on a standard laptop (Apple M3 Max), whereas training without the surrogate (using full simulations) would have required over 8,300 hours.
  • Spectral Adaptation: Analysis of the resulting group structures showed that the RL agents successfully adapted energy bounds to the specific neutron spectrum. For the fast-spectrum Godiva problem, the RL models concentrated bounds in the fast-energy region, whereas standard LANL30 structures placed more resolution in resonance/thermal regions less relevant to the specific problem.

Significance
The paper demonstrates that reinforcement learning, when coupled with surrogate modeling, offers a flexible and computationally efficient alternative to traditional group structure optimization techniques. The method successfully avoids local minima traps common in greedy algorithms and reduces the computational burden of optimization by orders of magnitude. By learning to remove bounds from a high-fidelity grid, the approach generates problem-specific group structures that outperform generalized standard grids (LANL30/70) while maintaining the ability to generalize across different material configurations and starting conditions without retraining. The authors note that future work could expand the action space to include adding or perturbing bounds and further refine surrogate resolution to improve performance.