It is mathematically correct to use a Policy Gradient method for 1-step trajectories?
The fundamental idea behind policy gradient is just to maximise the return averaged across all probably trajectories, i.e
$$\begin{align} J(\theta) &amp;= E[\sum\limits_{t=1}^{\tau}r(s_t,a_t)]\\
&amp;=E_{\tau\sim p(\tau)}[R(\tau)]
\end{align}$$
Where $\tau$ represents the probability of selecting a particular trajectory, if the trajectories all have fixed length then $\tau$ only has non-zero probability for trajectories of the specified length which in this case is 1.
The REINFORCE algorithm takes this expression and with some simplifications (causality to improve variance) and manipulation (log trick) obtains the gradient, pretty much as simple as that.

Intuition of algorithm in paper
In the algorithm they denote the $\pi_\theta$ (which is typically reserved for describing the policy) as the probability of the selection vector. By considering instead that the function $h_{\theta}$ as the policy instead I think it can be seen that they are actually averaging over multiple trajectories.
So we instead think of the data points as state-action pairs that we pass into $h_\theta$ to get the probability of selecting said action for a given state. These probabilities then dictates whether we choose the action or not. An alternative way to interpret this as if we imagine the action space for each state as binary then we can think of the &quot;other action&quot; as not impacting the predictor.
The gradient used for updating the parameters associated with the policy uses the log probability of $\pi_\theta$ which if we expand it, expressing it using $h_\theta$, (as done on page 4)  we can see it's the sum of the log probabilities of selecting (or not selecting) each data point.
By considering the return for each data point constant (each data point has the same loss incurred by predictor model on the validation set) and absorbing the average over batch size into the step size $\beta$ it could be interpreted as an average

Stability issues
RL is plagued with stability issues, be it selection of hyper parameters, random seeding etc. It's hard to often pinpoint why results aren't exact but as long as you get something in a similar ball park i'd say thats pretty good going