Small fix to AI