Implement the Flow Matching loss

[saleml]

Then try it on the HyperGrid environment (make the loss accessible via the configs for train.py). Make sure to reproduce the results reported in Trajectory Balance: Improved Credit Assignment in GFlowNets.

[danjenson]

Here is an implementation that may help, although my abstractions are slightly different (i.e. I have a BatchFlowStep class):
https://github.com/ThrunGroup/gfn/blob/main/gfn/losses.py