Cleaned up the policy function which fixes #4.