SB3: PPO for Knights-Archers-Zombies

This tutorial shows how to train agents using Proximal Policy Optimization (PPO) on the Knights-Archers-Zombies environment (AEC).

We use SuperSuit to create vectorized environments, leveraging multithreading to speed up training (see SB3’s vector environments documentation).

After training and evaluation, this script will launch a demo game using human rendering. Trained models are saved and loaded from disk (see SB3’s model saving documentation).

If the observation space is visual (vector_state=False in env_kwargs), we pre-process using color reduction, resizing, and frame stacking, and use a CNN policy.

Note

This environment has a visual (3-dimensional) observation space, so we use a CNN feature extractor.

Note

This environment allows agents to spawn and die, so it requires using SuperSuit’s Black Death wrapper, which provides blank observations to dead agents rather than removing them from the environment.

Environment Setup

To follow this tutorial, you will need to install the dependencies shown below. It is recommended to use a newly-created virtual environment to avoid dependency conflicts.

pettingzoo[butterfly]>=1.27.0
stable-baselines3>=2.0.0
supersuit>=3.9.0
pillow>=8.0.1

Code

The following code should run without any issues. The comments are designed to help you understand how to use PettingZoo with SB3. If you have any questions, please feel free to ask in the Discord server.

Training and Evaluation

"""Uses Stable-Baselines3 to train agents in the Knights-Archers-Zombies environment using SuperSuit vector envs.

This environment requires using SuperSuit's Black Death wrapper, to handle agent death.

For more information, see https://stable-baselines3.readthedocs.io/en/master/modules/ppo.html

Author: Elliot (https://github.com/elliottower)
"""

from __future__ import annotations

import glob
import os
import time

import supersuit as ss
from stable_baselines3 import PPO
from stable_baselines3.ppo import CnnPolicy, MlpPolicy

from pettingzoo import make


def train(env_id, steps: int = 10_000, seed: int | None = 0, **env_kwargs):
    # Train a single model to play as each agent in an AEC environment
    env = make("parallel", env_id, **env_kwargs)

    # Add black death wrapper so the number of agents stays constant
    # MarkovVectorEnv does not support environments with varying numbers of active agents unless black_death is set to True
    env = ss.black_death_v3(env)

    # Pre-process using SuperSuit
    visual_observation = not env.unwrapped.vector_state
    if visual_observation:
        # If the observation space is visual, reduce the color channels, resize from 512px to 84px, and apply frame stacking
        env = ss.color_reduction_v0(env, mode="B")
        env = ss.resize_v1(env, x_size=84, y_size=84)
        env = ss.frame_stack_v1(env, 3)

    env.reset(seed=seed)
    env_name = str(env.metadata.get("name", "pettingzoo_env"))

    print(f"Starting training on {env_name}.")

    env = ss.pettingzoo_env_to_vec_env_v1(env)
    env = ss.concat_vec_envs_v1(env, 8, num_cpus=1, base_class="stable_baselines3")

    # Use a CNN policy if the observation space is visual
    model = PPO(
        CnnPolicy if visual_observation else MlpPolicy,
        env,
        verbose=3,
        batch_size=256,
    )

    model.learn(total_timesteps=steps)

    model.save(f"{env_name}_{time.strftime('%Y%m%d-%H%M%S')}")

    print("Model has been saved.")

    print(f"Finished training on {env_name}.")

    env.close()


def eval(env_id, num_games: int = 100, render_mode: str | None = None, **env_kwargs):
    # Evaluate a trained agent vs a random agent
    env = make("aec", env_id, render_mode=render_mode, **env_kwargs)

    # Pre-process using SuperSuit
    visual_observation = not env.unwrapped.vector_state
    if visual_observation:
        # If the observation space is visual, reduce the color channels, resize from 512px to 84px, and apply frame stacking
        env = ss.color_reduction_v0(env, mode="B")
        env = ss.resize_v1(env, x_size=84, y_size=84)
        env = ss.frame_stack_v1(env, 3)

    print(
        f"\nStarting evaluation on {env.metadata['name']!s} (num_games={num_games}, render_mode={render_mode})"
    )

    try:
        latest_policy = max(
            glob.glob(f"{env.metadata['name']}*.zip"), key=os.path.getctime
        )
    except ValueError:
        print("Policy not found.")
        exit(0)

    model = PPO.load(latest_policy)

    rewards = dict.fromkeys(env.possible_agents, 0)

    # Note: we evaluate here using an AEC environments, to allow for easy A/B testing against random policies
    # For example, we can see here that using a random agent for archer_0 results in less points than the trained agent
    for i in range(num_games):
        env.reset(seed=i)
        env.action_space(env.possible_agents[0]).seed(i)

        for agent in env.agent_iter():
            obs, reward, termination, truncation, info = env.last()

            for a in env.agents:
                rewards[a] += env.rewards[a]

            if termination or truncation:
                break
            if agent == env.possible_agents[0]:
                act = env.action_space(agent).sample()
            else:
                act = model.predict(obs, deterministic=True)[0]
            env.step(act)
    env.close()

    avg_reward = sum(rewards.values()) / len(rewards.values())
    avg_reward_per_agent = {
        agent: rewards[agent] / num_games for agent in env.possible_agents
    }
    print(f"Avg reward: {avg_reward}")
    print("Avg reward per agent, per game: ", avg_reward_per_agent)
    print("Full rewards: ", rewards)
    return avg_reward


if __name__ == "__main__":
    env_id = "butterfly/knights_archers_zombies-v11"

    # Set obs_method to "image" in order to use visual observations (significantly longer training time)
    env_kwargs = {"max_cycles": 100, "max_zombies": 4, "obs_method": "vector"}

    # Train a model (takes ~5 minutes on a laptop CPU)
    train(env_id, steps=81_920, seed=0, **env_kwargs)

    # Evaluate 10 games (takes ~10 seconds on a laptop CPU)
    eval(env_id, num_games=10, render_mode=None, **env_kwargs)

    # Watch 2 games (takes ~10 seconds on a laptop CPU)
    eval(env_id, num_games=2, render_mode="human", **env_kwargs)

Evolved Vector-Policy Demo

A lightweight vector-observation policy playing Knights-Archers-Zombies.

For a fast, reproducible policy baseline, the script below uses the vector-masked observation, runs a small grid search over interpretable aiming parameters, predicts arrow/zombie interception points, evaluates the tuned policy against a random baseline, and can render an episode as a GIF.

On the default demo configuration below (max_cycles=900, max_zombies=10), the tuned policy averaged 43.7 total reward over seeds 0-49, compared with 2.4 for seeded random actions. Pass --search to repeat the small parameter sweep before evaluation.

python tutorials/SB3/kaz/evolved_kaz_policy.py \
    --episodes 50 \
    --max-cycles 900 \
    --max-zombies 10 \
    --render-gif docs/tutorials/sb3/kaz_evolved_policy.gif \
    --gif-seed 1

"""Interpretable vector-observation policy for Knights-Archers-Zombies.

The policy uses a small reproducible parameter sweep, projectile interception
for the archers, and a close-range fallback for the knights. It can evaluate
the policy against seeded random actions and render an episode as a GIF.
"""

from __future__ import annotations

import argparse
import itertools
import math
from dataclasses import asdict, dataclass, replace
from pathlib import Path

import numpy as np
from PIL import Image

from pettingzoo.butterfly import knights_archers_zombies_v11
from pettingzoo.butterfly.knights_archers_zombies.src import constants as const


@dataclass(frozen=True)
class PolicyParams:
    """Tunable parameters for the deterministic KAZ policy."""

    archer_align_degrees: float = 6.6
    archer_lane_split: float = 0.70
    left_standby_degrees: float = 0.0
    right_standby_degrees: float = 35.0
    knight_track_distance: float = 260.0
    knight_attack_distance: float = 150.0
    knight_track_degrees: float = 18.0
    knight_attack_degrees: float = 25.0


def zombie_rows(observation: np.ndarray) -> list[np.ndarray]:
    """Return active zombie rows from a vector-masked observation."""
    return [row for row in observation if row[0] > 0.5]


def relative_pixels(row: np.ndarray) -> tuple[float, float]:
    """Convert an entity's normalized relative position to screen pixels."""
    return (
        float(row[7]) * const.SCREEN_WIDTH,
        float(row[8]) * const.SCREEN_HEIGHT,
    )


def normalized_direction(x: float, y: float) -> tuple[float, float]:
    """Return a unit direction, defaulting upward for a zero-length vector."""
    length = math.hypot(x, y)
    if length < 1e-9:
        return 0.0, -1.0
    return x / length, y / length


def intercept_direction(row: np.ndarray) -> tuple[float, float]:
    """Aim where a downward-moving zombie and an arrow should intersect."""
    x, y = relative_pixels(row)
    zombie_speed = float(const.ZOMBIE_Y_SPEED)
    arrow_speed = float(const.ARROW_SPEED)

    # Solve |(x, y) + (0, zombie_speed) * t| = arrow_speed * t.
    a = zombie_speed**2 - arrow_speed**2
    b = 2.0 * y * zombie_speed
    c = x**2 + y**2
    discriminant = max(b**2 - 4.0 * a * c, 0.0)
    roots = (
        (-b + math.sqrt(discriminant)) / (2.0 * a),
        (-b - math.sqrt(discriminant)) / (2.0 * a),
    )
    positive_roots = [root for root in roots if root > 0.0]
    intercept_time = (
        min(positive_roots) if positive_roots else math.sqrt(c) / arrow_speed
    )
    return normalized_direction(x, y + zombie_speed * intercept_time)


def turn_action(
    observation: np.ndarray,
    target_x: float,
    target_y: float,
    *,
    tolerance_degrees: float,
) -> int | None:
    """Return a turn action, or None when the agent is sufficiently aligned."""
    heading_x = float(observation[0][9])
    heading_y = float(observation[0][10])
    dot = heading_x * target_x + heading_y * target_y
    if dot >= math.cos(math.radians(tolerance_degrees)):
        return None

    cross = heading_x * target_y - heading_y * target_x
    return 3 if cross > 0.0 else 2


def archer_standby_action(
    observation: np.ndarray, archer_index: int, params: PolicyParams
) -> int:
    """Fan idle archers across the board so new targets need little turning."""
    angle_degrees = (
        params.left_standby_degrees
        if archer_index % 2 == 0
        else params.right_standby_degrees
    )
    angle = math.radians(angle_degrees)
    action = turn_action(
        observation,
        math.sin(angle),
        -math.cos(angle),
        tolerance_degrees=5.0,
    )
    return 5 if action is None else action


def select_archer_target(
    observation: np.ndarray,
    zombies: list[np.ndarray],
    archer_index: int,
    params: PolicyParams,
) -> np.ndarray:
    """Select an urgent target while giving one archer the far-right lane."""
    current_x = float(observation[0][7])
    primary_lane_is_left = archer_index % 2 == 0

    def target_key(row: np.ndarray) -> tuple[int, float, float]:
        world_x = current_x + float(row[7])
        outside_primary_lane = int(
            (world_x >= params.archer_lane_split) == primary_lane_is_left
        )
        relative_x, relative_y = relative_pixels(row)
        return outside_primary_lane, math.hypot(relative_x, relative_y), -relative_y

    return min(zombies, key=target_key)


def policy_action(
    observation: np.ndarray,
    agent: str,
    params: PolicyParams = PolicyParams(),
) -> int:
    """Map a vector-masked observation to a deterministic KAZ action."""
    zombies = zombie_rows(observation)

    if agent.startswith("archer"):
        archer_index = int(agent.rsplit("_", 1)[1])
        if not zombies:
            return archer_standby_action(observation, archer_index, params)

        target = select_archer_target(observation, zombies, archer_index, params)
        target_x, target_y = intercept_direction(target)
        action = turn_action(
            observation,
            target_x,
            target_y,
            tolerance_degrees=params.archer_align_degrees,
        )
        return 4 if action is None else action

    if not zombies:
        return 5

    target = min(zombies, key=lambda row: math.hypot(*relative_pixels(row)))
    relative_x, relative_y = relative_pixels(target)
    distance = math.hypot(relative_x, relative_y)
    if distance >= params.knight_track_distance:
        return 5

    target_x, target_y = normalized_direction(relative_x, relative_y)
    tolerance = (
        params.knight_attack_degrees
        if distance <= params.knight_attack_distance
        else params.knight_track_degrees
    )
    action = turn_action(
        observation,
        target_x,
        target_y,
        tolerance_degrees=tolerance,
    )
    if action is not None:
        return action
    return 4 if distance <= params.knight_attack_distance else 5


def run_episode(
    seed: int,
    params: PolicyParams = PolicyParams(),
    *,
    max_cycles: int = 900,
    max_zombies: int = 10,
    random_policy: bool = False,
    render: bool = False,
    frame_stride: int = 5,
) -> tuple[float, list[np.ndarray]]:
    """Run one KAZ episode and optionally capture RGB frames."""
    env = knights_archers_zombies_v11.env(
        render_mode="rgb_array" if render else None,
        max_cycles=max_cycles,
        max_zombies=max_zombies,
        max_arrows=10,
        obs_method="vector-masked",
    )
    env.reset(seed=seed)
    for agent_index, agent in enumerate(env.possible_agents):
        env.action_space(agent).seed(seed * 1009 + agent_index)

    total_reward = 0.0
    frames: list[np.ndarray] = []
    step = 0
    for agent in env.agent_iter():
        observation, reward, termination, truncation, _ = env.last()
        total_reward += reward

        if termination or truncation:
            action = None
        elif random_policy:
            action = env.action_space(agent).sample()
        else:
            action = policy_action(observation, agent, params)

        env.step(action)
        if render and step % frame_stride == 0:
            frame = env.render()
            if frame is not None:
                frames.append(frame)
        step += 1

    env.close()
    return total_reward, frames


def evaluate(
    seeds: range,
    params: PolicyParams = PolicyParams(),
    *,
    max_cycles: int,
    max_zombies: int,
    random_policy: bool = False,
) -> list[float]:
    """Evaluate a policy on every seed in a range."""
    return [
        run_episode(
            seed,
            params,
            max_cycles=max_cycles,
            max_zombies=max_zombies,
            random_policy=random_policy,
        )[0]
        for seed in seeds
    ]


def search_params(
    *, seeds: range, max_cycles: int, max_zombies: int
) -> tuple[PolicyParams, list[float]]:
    """Run a small grid search over interpretable archer parameters."""
    defaults = PolicyParams()
    best_params = defaults
    best_scores: list[float] = []
    best_mean = float("-inf")

    for align, lane_split, left_standby, right_standby in itertools.product(
        [6.0, 6.6, 7.0],
        [0.65, 0.70, 0.75],
        [-5.0, 0.0],
        [30.0, 35.0],
    ):
        params = replace(
            defaults,
            archer_align_degrees=align,
            archer_lane_split=lane_split,
            left_standby_degrees=left_standby,
            right_standby_degrees=right_standby,
        )
        scores = evaluate(
            seeds,
            params,
            max_cycles=max_cycles,
            max_zombies=max_zombies,
        )
        mean_score = sum(scores) / len(scores)
        if mean_score > best_mean:
            best_mean = mean_score
            best_params = params
            best_scores = scores

    return best_params, best_scores


def save_gif(frames: list[np.ndarray], output_path: Path, *, duration_ms: int) -> None:
    """Save captured RGB frames as a looping GIF."""
    if not frames:
        raise ValueError("No frames were captured.")

    images = [Image.fromarray(frame) for frame in frames]
    output_path.parent.mkdir(parents=True, exist_ok=True)
    images[0].save(
        output_path,
        save_all=True,
        append_images=images[1:],
        duration=duration_ms,
        loop=0,
    )


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--episodes", type=int, default=10)
    parser.add_argument("--seed-start", type=int, default=0)
    parser.add_argument("--max-cycles", type=int, default=900)
    parser.add_argument("--max-zombies", type=int, default=10)
    parser.add_argument("--search", action="store_true")
    parser.add_argument("--render-gif", type=Path)
    parser.add_argument("--gif-seed", type=int, default=0)
    parser.add_argument("--gif-duration-ms", type=int, default=80)
    args = parser.parse_args()
    if args.episodes <= 0:
        parser.error("--episodes must be positive")
    return args


def main() -> None:
    """Evaluate the tuned policy and optionally search or render it."""
    args = parse_args()
    seeds = range(args.seed_start, args.seed_start + args.episodes)

    params = PolicyParams()
    if args.search:
        search_seeds = range(args.seed_start, args.seed_start + min(args.episodes, 5))
        params, search_scores = search_params(
            seeds=search_seeds,
            max_cycles=args.max_cycles,
            max_zombies=args.max_zombies,
        )
        print("Best searched params:", asdict(params))
        print("Search scores:", search_scores)

    policy_scores = evaluate(
        seeds,
        params,
        max_cycles=args.max_cycles,
        max_zombies=args.max_zombies,
    )
    random_scores = evaluate(
        seeds,
        params,
        max_cycles=args.max_cycles,
        max_zombies=args.max_zombies,
        random_policy=True,
    )
    print("Policy params:", asdict(params))
    print(f"Policy scores: {policy_scores}")
    print(f"Random scores: {random_scores}")
    print(f"Policy mean: {sum(policy_scores) / len(policy_scores):.2f}")
    print(f"Random mean: {sum(random_scores) / len(random_scores):.2f}")

    if args.render_gif is not None:
        reward, frames = run_episode(
            args.gif_seed,
            params,
            max_cycles=args.max_cycles,
            max_zombies=args.max_zombies,
            render=True,
        )
        save_gif(frames, args.render_gif, duration_ms=args.gif_duration_ms)
        print(f"Rendered seed {args.gif_seed} with reward {reward:.0f}")
        print(f"Wrote {args.render_gif}")


if __name__ == "__main__":
    main()