I realised I had some logical errors in my step function which I have now fixed:

Previously, the function checked whether the rocket landed, crashed or was in the air then it implemented changes to the rocket’s y coordinates via the velocity depending on the action taken

This order of working is wrong and should have been reversed. The check for landing, crashing or flying should be implemented after the rocket’s y coordinates have been updated in accordance to the action taken.

The AI has now overcome the problem, hence, I must increase the complexity.

I have done this by introducing acceleration due to gravity and a fuel system that prevents the rocket from acceleration if it runs out of fuel. This forces the AI to think carefully about using its thrust rather than exploiting the previous infinite fuel.

With this implementation, the rocket landing is much more realistic and smoother.