def test_mask_from_labels(): # Test the create_mask helper separately space = BoxSpace(-1, 1, shape=5, labels=['w', 'o', 'r', 'l', 'd']) indices = ['w', 'o'] mask = space.create_mask(indices) assert list(mask) == [1, 1, 0, 0, 0]
def test_mask_from_indices(): # Test the create_mask helper separately space = BoxSpace(-1, 1, shape=5) indices = [1, 4] mask = space.create_mask(indices) assert list(mask) == [0, 1, 0, 0, 1]