Tensorflow / Deepmind: how do I take actions from observations for math algorithms related to proofs?
I would recommend you to have a look at OpenAI gym.
It will be easier for you to implement a new environment describing your problem. Regarding the agents, there are many implementations around that are compatible with OpenAI.