def test_wrapper():
    """
    **Overview**:
        Test two types wrappers and check output obs shape.
    """
    # Test environment with stacked frames.
    env = wrapped_mario_env()
    obs = env.reset()
    assert obs.shape == (4, 84, 84)
    # Test environment with optical flow.
    env = wrapped_mario_env_optical()
    obs = env.reset()
    assert obs.shape == (3, 84, 84)

    **Overview**:
        Test two types wrappers and check output obs shape.
    
chapter3_obs/mario_wrapper.py