I hope I wont be in trouble for over-promoting my work… I think this is the solution you are looking for.
I haven’t ported this to Python yet, this is still an open problem.
However, to create a learning agent utilising Phasm is a completely reasonable task.
https://chatgpt.com/share/67cbe59b-f7a8-8003-915c-1414d27893ac