Author of the RWA model here. I saw your post, and I wanted to let you know a flaw has been discovered in my code. The flaw deals with the numerical stability of the RWA model. If left uncorrected it prevents the model from forming long-term memories. Once you fix your code you may discover the issue you are having go away. Maybe so, maybe not.
I just posted the corrected code on my repo (https://github.com/jostmey/rwa 56). I owe a special thanks to some dude named Alex Nichol (a.k.a. unixpickle) for finding the bug. I am re-running all my results. So far the corrected model appears to run at least as well as before