Thanks! We did try feeding both the "cell" of the LSTM in addition to the hidden state of the LSTM and the latent vector into the controller, and this works better. We discussed this in the Appendix section.

I would still encourage you to pursue your idea, since there are still lots of limitations in this model (discussed in the paper), and a lot of work remains to be done to solve more difficult problems.