Custom Covariance Function for Gaussian Process

Thanks for the response.
The idea is to generate a mask based on the reward_trials. If the reward_trials == 1 at i, we want to set the entries after and below the diagonal (i,i) to be 1. Applying this mask to the original kernel will produce the second kernel matrix shown in the figure.