from pettingzoo.butterfly import knights_archers_zombies_v4
env = knights_archers_zombies_v4.env()
env.reset()
for agent in env.agent_iter():
    observation, reward, done, info = env.last()
    action = policy(observation, agent)
    env.step(action)