Improve ai behavior