Optimizing a superconducting radio-frequency gun using deep reinforcement learning

Superconducting photoelectron injectors are promising for generating highly brilliant pulsed electron beams with high repetition rates and low emittances. Experiments such as ultrafast electron diffraction, experiments at the Terahertz scale, and energy recovery linac applications require such properties. However, optimizing the beam properties is challenging due to the high number of possible machine parameter combinations. This article shows the successful automated optimization of beam properties utilizing an already existing simulation model. To reduce the required computation time, we replace the costly simulation with a faster approximation with a neural network. For optimization, we propose a reinforcement learning approach leveraging the simple computation of the derivative of the approximation. We prove that our approach outperforms standard optimization methods for the required function evaluations given a defined minimum accuracy.

Figure 1

Schematic view of the electron gun; the input variables are the parameters of the cathode, solenoids, and gun cavity. The output beam parameters can also be measured in the real device.

Figure 2

Schematic representation of the rl cycle; The agent illustrates our RL algorithm. It can apply an action $a$ to the environment. The environment depicts in our case, the simulation approximation which can calculate its new state $s$ and the reward $R_l$. It gives this information to the RL agent to decide the next action.

Figure 3

Basic processing pipeline; this article focuses on the optimization part.

Figure 4

Schema of the parameters’ role within the learning loop. The blue arrows indicate the transitions within the pretrained surrogate model. The green elements indicate the training and decision-making by the RL agent. First, the RL agent determines how to choose the action variables using his policy. The state and integration variables are chosen randomly, but we can later measure them in the real experiment. Together with the state and integration variables, these actions get evaluated by the pretrained surrogate model. It returns the optimization targets, which serve as a reward for the policy, from which the RL agent can then improve its policy. Now the learning cycle repeats.

Figure 5

Schematic plot of the neural network we use for learning the policy. We scaled the heights of the layers logarithmically for visualization purposes. The numbers indicate the layer sizes, input and output dimensions.

Figure 6

Comparison of the evaluated RL policy $\mu (s)$ with the best reward achieved by the different optimization methods after 1000 reward function evaluations; Because the local optimization algorithms are stochastic, i.e., depend on random variables, we run the optimization algorithms 800 times to get a reliable amount of statistics. We plot the histograms with respect to the optimal rewards of our RL agent. The intensity shows how many runs of the stochastic optimizations have achieved a particular optimal reward. In case one method is performing equally as RL, only the bisecting line (highlighted in red) will be visible. If there is intensity below this line, the optimization method could not achieve a similar or better reward than the RL policy $\mu (s)$ within the allowed reward function evaluations. Note the logarithmic scale on both axes.

Figure 7

Comparison of the evaluated RL policy $\mu (s)$ with local optimization algorithms with respect to achieved reward at different evaluation counts. A larger reward is better. We show the mean of 800 repetitions over 1000 evaluation steps. The standard deviation is plotted semitransparent. The RL policy is plotted as a dashed line to highlight that it is only a fixed value since, after completing training, it only requires one evaluation of the reward function.

Figure 8

Average of the Jacobian matrix of 100,000 state values and their chosen actions; A high value in the color scale means the average of the partial derivatives is high. It turns out that great changes in the laser horizontal position and vertical position lead to great changes in the respective solenoid position in the policy of our RL agent.

Figure 9

Standard deviations of the Jacobian matrix of 100,000 state values and their chosen actions; A high value in the color scale means the standard deviation of the partial derivatives is high. It shows that the derivatives of the solenoid position have high fluctuations both vertically and horizontally with respect to the cathode position.

Figure 10

Random action values ordered by laser horizontal positions; The standard deviation of the input values is $\sigma =0.2$. This plot shows the mean of 500 laser horizontal positions and 1000 repetitions. The variance is plotted semitransparent. The plot shows that with increasing laser horizontal position, also the solenoid horizontal position is increased. The solenoid angle in the $y$ axis is also moved slightly. However, the solenoid vertical position and solenoid angle $x$ axis are almost not altered.

Figure 11

The training process of the policy. Since the training is stochastic, we repeat it 500 times to get sufficient statistical evidence. The opaque graph shows the mean of all repetitions and the standard deviation is plotted semitransparent. The achieved reward improves only slightly after 20,000 epochs.