Tutorial: Testing Your Environment

Introduction

Now that our environment is complete, we can test it to make sure it works as intended. PettingZoo has a built-in testing suite that can be used to test your environment.

After the API tests pass, register the environment with pettingzoo.register so that it can be created with make, then confirm it appears in parallel_registry (or aec_registry for AEC environments). See Basic Usage for more on the registry-related APIs.

Code

Note: This code can be added to the bottom of the same file, without using any imports, but it is best practice to keep tests in a separate file, and use modular imports, as shown below..

Relative importing is used for simplicity, and assumes your custom environment is in the same directory. If your test is in another location (e.g., a root-level /test/ directory), it is recommended to import using absolute path.

/custom-environment/env/custom_environment.py

from tutorial2_adding_game_logic import CustomEnvironment
from tutorial3_action_masking import CustomActionMaskedEnvironment

from pettingzoo import make, parallel_registry, register
from pettingzoo.test import parallel_api_test

if __name__ == "__main__":
    # Sanity check
    env = CustomEnvironment()
    parallel_api_test(env, num_cycles=1_000_000)

    env = CustomActionMaskedEnvironment()
    parallel_api_test(env, num_cycles=1_000_000)

    # Register custom environments so they can be created with make()
    register("parallel", "custom/custom_environment-v0", CustomEnvironment)
    register(
        "parallel",
        "custom/custom_action_masked_environment-v0",
        CustomActionMaskedEnvironment,
    )

    # Confirm the environments are available in the parallel registry
    assert "custom/custom_environment-v0" in parallel_registry
    assert "custom/custom_action_masked_environment-v0" in parallel_registry

    # Environments can now be created via make, just like built-in ones
    env = make("parallel", "custom/custom_environment-v0")
    parallel_api_test(env, num_cycles=1_000_000)

    env = make("parallel", "custom/custom_action_masked_environment-v0")
    parallel_api_test(env, num_cycles=1_000_000)