This is called "distance metric learning" or "learning distance metrics". You can find many papers in the literature on this topic.
There are many algorithms for learning distance metrics. Some of them allow you to specify quadruplets $(x_1,x_2,x_3,x_4)$ where we are promised that $d(x_1,x_2) < d(x_3,x_4)$, and the task is to learn a distance matric $d$ that is consistent with this training set. You can then use such an algorithm for your task, by sampling $x_1$ randomly from your training set, sampling $x_2$ from one of the $K$ nearest neighbors to $x_1$, setting $x_3=x_1$, and sampling $x_4$ from one of the other $N-K-1$ non-neighbors.
One approach is to learn a Mahalanobis distance, i.e., a distance metric of the form $d(x,x') = \|Lx-Lx'\|_2$. This can be equivalently formulated as $d(x,x') = \sqrt{(x-x')^\top M(x-x')}$ where $M=L^\top L$. You can formulate the learning task as an optimization problem, and then use standard optimization methods to learn the matrix $L$ (or $M$). This learns a linear distance metric.
There are other approaches that attempt to learn a more complex, nonlinear distance metric using a neural network, e.g., $d(x,x') = \|N(x)-N(x')\|_2$, where $N$ is a neural network (thus, measuring distance using a Siamese network). Standard literature on training neural networks to measure image similarity describe a number of ways to learn such a network, and it could be applied to your situation as well by sampling triplets appropriately.
You might be interested in the metric-learn package for Python.