Single-shot policy

Straight from what the sensors see to the knob settings, trained without a single optimal action.

A lot of accelerator tuning is one job done over and over. The beam coming into a transfer line shifts a little, and the best corrector settings shift with it. So we run the optimizer again, as if we'd never seen the problem. You could train a network to jump straight to the settings. But every training example would need its answer, and each answer is one more optimizer run (which may not have found the best one). My idea was to leave the answers out. Put a reward model you can differentiate after the policy network, and let the gradient of the reward do the teaching. I called it a single-shot policy. We tried it in simulation, and on one laser bench at LBNL. The write-up is LBNL-led (Dan Wang et al., I'm fifth of seven) and isn't published yet.

One step

A reinforcement-learning agent acts, looks at what happened, acts again, and adds up its reward over the episode. A lot of tuning problems don't need the second look. Once the monitors tell you the state, the state already decides the best action (the paper calls these deterministic problems). So the episode can be a single action, scored on its own. That's the single shot.

Two boxed diagrams, one above the other. Top, labelled RL Loop: State and Action boxes with arrows going both ways, and the line State init, Action 1, Action 2, and so on, State opt. A green arrow points down to the bottom one, labelled Single Shot: one arrow from State to Action star, and the line State init, Action star, State opt.
My original sketch, the reinforcement-learning loop on top and the single shot below.

Training through the reward

There are two networks. R-Net (R) takes a state and an action and predicts the reward. P-Net (P) takes a state and proposes an action. Feed P-Net's action into R-Net, and the predicted reward can be differentiated all the way back into P-Net's weights.

r=R(𝐒,P(𝐒))

𝐒 is the state, P the policy network and R the reward model, so r is the reward predicted for the action the policy itself proposed.

So the loss is just minus the mean reward over a batch of states. Nowhere does it ask what the right action was.

Training in two stages, then use. Stage one: random states and actions go into R-Net, which predicts the reward and is fitted to the measured one; this stage is skipped when the simulation is differentiable. Stage two: a state goes into P-Net, which proposes an action; the action goes into the frozen R-Net, which predicts the reward; the loss is minus the mean reward, and its gradient updates P-Net only. On the machine: the sensors give the state, P-Net gives the action in one forward pass, and the controller applies it. First, the reward model Then the policy, against the frozen R-Net On the machine skipped when the simulation can be differentiated state, actionR-Netreward stateP-NetactionR-Netreward stateP-Netaction random pairsbeing fittedvs the measured one a batch of thembeing trainedproposedfrozenloss = −mean from the sensorsto the controller its gradient updates P-Net only one forward pass, no search
We fit the reward model first and train the policy against it frozen, and after that the policy runs alone.

The transfer line

The first test was the line that takes electrons from the booster into SPEAR3 at SSRL, in simulation. (RCDS-S had chased a drift on the same line in simulation.) The policy set four correctors near the start (two horizontal, two vertical) and read x and y at all nine BPMs. What it wanted was injection efficiency at the far end. It couldn't see the incoming beam. We randomized its position and angle.

Schematic of the transfer line from left to right: a random upstream beam, then a row of horizontal correctors (blue, h1 to h9), vertical correctors (red, v1 to v7) and BPMs (yellow, b1 to b9), ending at injection efficiency. Dashed boxes mark h1, v1, h5 and v2 near the start.
The four correctors in dashed boxes are the policy's knobs, and b1 to b9 are the BPMs it reads.

Accelerator Toolbox, the MATLAB code we simulated it in, can't be differentiated, so this case needed an R-Net, fitted on a few thousand Latin-hypercube samples. The test was 1,000 incoming beams neither network had seen. Before the step, efficiency was often close to zero. After it, the mean was close to the best possible (about 0.904).

Histogram of normalized injection efficiency for the 1,000 test beams. The initial values, in blue, spread across the whole range from 0 to 1 with a tall spike near zero. The values after the policy's step, in orange, sit in a narrow peak just under 1.
Normalized injection efficiency for the 1,000 test beams, before the step (blue, with a tall pile near zero) and after it (orange, packed near 1).

Combining laser beams

Coherent beam combining is the laser version of the same kind of problem. Several beams overlap on a diffractive optic, and how much power lands in the centre spot depends on their relative phases. Those drift (heat alone will do it). The camera only shows you the pattern they make. In our simulation eight beams sat in a 3×3 grid with no centre beam. The policy saw a 5×5 patch of the camera and set eight phases, and the reward was the centre pixel, 64 at most. We wrote the simulation in PyTorch, so it could be differentiated and needed no R-Net.

Trained on only 100 random starts, the policy already landed just under 64. More starts helped. With 3,000 it got closer still, and the spread tightened a lot.

Two 5×5 camera patterns side by side with an arrow between them: on the left a random start, light scattered over several pixels; on the right the target, one bright centre pixel.
A random start and the pattern we want, on the 5×5 patch.
Histogram of centre intensity: the random starts in blue spread from about 54 to 63; the policies trained on 100 starts (orange) and on 3,000 (green) sit just under 64, the green one narrower.
Before (blue) and after, with the policy trained on 100 starts (orange) or 3,000 (green).

Then 9×9, 81 beams, the size LBNL's group is working toward. No 9×9 setup was available yet, so this only checked that the method scales. The policy read a 17×17 patch and set 81 phases. We trained it on 30,000 starts, none more than 30 degrees out of phase. It raised the mean centre intensity from about 5,990 to about 6,510, and the paper says outright that it did worse than the small case. Why it did worse is still just a note of mine in the draft.

Two 17×17 camera patterns side by side with an arrow between them: on the left a random start with many scattered bright pixels; on the right the target, a single bright spot in the centre.
The same picture for 81 beams, on a 17×17 patch.
Histogram of centre intensity for the 9×9 case: the random starts in blue form a broad bump centred near 6,000; the policy's results in orange form a narrow peak near 6,510.
Before (blue) and after (orange), for 81 beams.

The bench at LBNL

The hardware tests were on LBNL's eight-way combining bench, and the trace below is from 7 September 2023. The bench belongs to LBNL's laser-control group, and so does DDRM, the method they had for keeping it combined. It was built for fibre amplifiers that weren't in yet. So for this test one CW laser was split eight ways (same 3×3 grid, no centre) and piezo mirrors set the phases.

No differentiable model of the real bench exists, so this time R-Net had to learn from data. We ran SPGD on the bench (stochastic parallel gradient descent, which jiggles all the phases and follows whatever raises the centre). We logged 20,000 camera patterns, each with the phase step that followed. R-Net trained on them in under a minute. P-Net needed no live data (we generated its patterns) and was done within a minute too. The fiddly part was the scale on P-Net's output. Too large and the policy wanders where R-Net has never been, too small and good answers are out of reach. We settled on 1.5 after a few tries.

Centre-spot power over about ten minutes on 7 September 2023, with a red dashed reference line near the top. Labelled across the plot: single-shot policy (flat on the line), free drifting (dips and jitter below it), detuning in a red band (power near zero), single-shot policy again (jumps straight back to the line), free drifting, and LBNL's DDRM at the end (flat, just under the line).
About ten minutes on the bench, with the policy holding, free drift, a deliberate detune (red), the policy back on, drift again, and LBNL's DDRM at the end.

With the policy on, the centre power held steady. Paused, it sagged and jittered. Then we threw the phases off on purpose, switched the policy back on, and the power came straight back to where it had been. Against DDRM at the end, the manuscript reads it as more combined power with similar, slightly worse stability. The draft has no numbers for this run yet, only the trace.

The unfinished paper

Everything on the SLAC side is simulation, and nothing was deployed. The draft has no timing against a search, so there's no speed-up to report. The paper hopes the policy can ride out drift without retraining. No result shows that yet. The draft isn't finished or submitted, and so far there's only an abstract, at IPAC'24.

Numbers and figures are from the draft manuscript, "Simplified reinforcement learning enabled deterministic feedback controls in lasers and accelerators: experimental demonstration on laser combining" (Wang, Du, Zhou, Wilcox, Zhang, Huang and Ratner, not yet published), shown here with Dan Wang's permission. The one-step sketch is mine.