Back to Repositories

Testing Observation Space Flattening Functionality in OpenAI Gym

This test suite validates the FlattenObservation wrapper functionality in OpenAI Gym, focusing on transforming multi-dimensional observation spaces into flattened array representations. The tests ensure proper handling of observation space conversion and information preservation during environment interactions.

Test Coverage Overview

The test coverage focuses on verifying the FlattenObservation wrapper’s core functionality with the Blackjack-v1 environment.

Validates observation space transformation from Tuple to Box space
Verifies preservation of observation data during reset operations
Tests proper type maintenance for observation and info dictionaries
Ensures compatibility with discrete space components

Implementation Analysis

The testing approach utilizes pytest’s parametrize decorator to enable flexible environment testing configurations. The implementation validates both the wrapped and unwrapped environment states, ensuring proper space conversion while maintaining data integrity.

The test leverages numpy arrays and Gym’s space definitions to verify correct dimensionality and type preservation in the flattened observation space.

Technical Details

Testing Framework: pytest
Key Dependencies: numpy, gym
Environment: Blackjack-v1
Wrapper: FlattenObservation
Space Types: spaces.Tuple, spaces.Discrete, spaces.Box
Data Types: np.int64

Best Practices Demonstrated

The test implementation showcases several testing best practices for wrapper validation.

Explicit space definition and verification
Comprehensive state checking before and after wrapping
Type safety validation for critical components
Parameterized test structure for extensibility
Clear separation of environment setup and assertion logic

openai/gym

tests/wrappers/test_flatten_observation.py

            
import numpy as np
import pytest

import gym
from gym import spaces
from gym.wrappers import FlattenObservation


@pytest.mark.parametrize("env_id", ["Blackjack-v1"])
def test_flatten_observation(env_id):
    env = gym.make(env_id, disable_env_checker=True)
    wrapped_env = FlattenObservation(env)

    obs, info = env.reset()
    wrapped_obs, wrapped_obs_info = wrapped_env.reset()

    space = spaces.Tuple((spaces.Discrete(32), spaces.Discrete(11), spaces.Discrete(2)))
    wrapped_space = spaces.Box(0, 1, [32 + 11 + 2], dtype=np.int64)

    assert space.contains(obs)
    assert wrapped_space.contains(wrapped_obs)
    assert isinstance(info, dict)
    assert isinstance(wrapped_obs_info, dict)