<commit_before>import numpy as np
import pytest

from imitation.util.buffer import Buffer, ReplayBuffer


@pytest.mark.parametrize("capacity", [10, 30, 60])
@pytest.mark.parametrize("chunk_len", [1, 2, 4, 9])
@pytest.mark.parametrize("sample_shape", [(), (1, 2), (5, 4, 4)])
def test_buffer(capacity, chunk_len, sample_shape):
  """Builds a Buffer with the provided `capacity` and insert `capacity * 3`
  samples into the buffer in chunks of shape `(chunk_len,) + sample_shape`.

  We always insert the same chunk, an array containing only 66.6.

  * `len(buffer)` should increase until we reach capacity.
  * `buffer._idx` should loop between 0 and `capacity - 1`.
  * After every insertion, samples should only contain 66.6.
  * Mutating the inserted chunk shouldn't mutate the buffer.
  """
  buf = Buffer(capacity, sample_shape=sample_shape, dtype=float)
  data = np.full((chunk_len,) + sample_shape, 66.6)

  for i in range(0, capacity*3, chunk_len):
    assert len(buf) == min(i, capacity)
    assert buf._idx == i % capacity
    buf.store(data)
    samples = buf.sample(100)
    assert samples.shape == (100,) + sample_shape
    assert np.all(samples == 66.6)

  # Confirm that buffer is not mutable from inserted sample.
  data[:] = np.nan
  samples = buf.sample(100)
  assert samples.shape == (100,) + sample_shape
  assert np.all(samples == 66.6)


@pytest.mark.parametrize("capacity", [30, 60])
@pytest.mark.parametrize("chunk_len", [1, 4, 9])
@pytest.mark.parametrize("obs_shape", [(), (1, 2)])
@pytest.mark.parametrize("act_shape", [(), (5, 4, 4)])
@pytest.mark.parametrize("dtype", [np.int, np.bool, np.float32])
def test_replay_buffer(capacity, chunk_len, obs_shape, act_shape, dtype):
  """Builds a ReplayBuffer with the provided `capacity` and inserts
  `capacity * 3` observation-action-observation samples into the buffer in
  chunks of length `chunk_len`.

  All chunks are of the appopriate observation or action shape, and contain
  the value fill_val.

  `len(buffer)` should increase until we reach capacity.
  `buffer._idx` should loop between 0 and `capacity - 1`.
  After every insertion, samples should only contain 66.6.
  """
  buf = ReplayBuffer(capacity, obs_shape=obs_shape, act_shape=act_shape,
                     obs_dtype=dtype, act_dtype=dtype)
  old_obs_fill_val, act_fill_val, new_obs_fill_val = (
      np.array(x, dtype=dtype) for x in [0, 1, 2])
  old_obs_data = np.full(
      (chunk_len,) + obs_shape, old_obs_fill_val, dtype=dtype)
  act_data = np.full((chunk_len,) + act_shape, act_fill_val, dtype=dtype)
  new_obs_data = np.full(
      (chunk_len,) + obs_shape, new_obs_fill_val, dtype=dtype)

  for i in range(0, capacity*3, chunk_len):
    assert len(buf) == min(i, capacity)
    for b in [buf._old_obs_buffer, buf._act_buffer, buf._new_obs_buffer]:
      assert b._idx == i % capacity

    buf.store(old_obs_data, act_data, new_obs_data)

    old_obs, acts, new_obs = buf.sample(100)
    assert old_obs.shape == new_obs.shape == (100,) + obs_shape
    assert acts.shape == (100,) + act_shape

    assert np.all(old_obs == old_obs_fill_val)
    assert np.all(acts == act_fill_val)
    assert np.all(new_obs == new_obs_fill_val)

    assert old_obs.dtype == dtype
    assert acts.dtype == dtype
    assert new_obs.dtype == dtype


@pytest.mark.parametrize("sample_shape", [(), (1,), (5, 2)])
def test_buffer_store_errors(sample_shape):
  capacity = 11
  dtype = "float32"

  def buf():
      return Buffer(capacity, sample_shape, dtype)

  # `data` is empty.
  b = buf()
  with pytest.raises(ValueError):
      b.store(np.empty((0,) + sample_shape, dtype=dtype))

  # `data` has too many samples.
  b = buf()
  with pytest.raises(ValueError):
      b.store(np.empty((capacity + 1,) + sample_shape, dtype=dtype))

  # `data` has the wrong sample shape.
  b = buf()
  with pytest.raises(ValueError):
      b.store(np.empty((1, 3, 3, 3, 3), dtype=dtype))


def test_buffer_sample_errors():
  b = Buffer(10, (2, 1), dtype=bool)
  with pytest.raises(ValueError):
      b.sample(5)
<commit_msg>Add test case catching bug<commit_after>import numpy as np
import pytest

from imitation.util.buffer import Buffer, ReplayBuffer


def _fill_chunk(start, chunk_len, sample_shape, dtype=np.float):
  fill_vals = np.arange(start, start + chunk_len, dtype=dtype)
  fill_vals = np.reshape(fill_vals, (-1, ) + (1,) * len(sample_shape))
  chunk = np.tile(fill_vals, (1,) + sample_shape)
  return chunk


def _get_fill_from_chunk(chunk):
  chunk_len, *sample_shape = chunk
  sample_size = max(1, np.prod(sample_shape))
  return chunk.flatten()[::sample_size]


def _check_bound(end, capacity, samples, offset=0):
  start = max(0, end - capacity)
  assert np.all(start + offset <= samples), "samples violate lower bound"
  assert np.all(samples <= end + offset), "samples violate upper bound"


@pytest.mark.parametrize("capacity", [10, 30, 60])
@pytest.mark.parametrize("chunk_len", [1, 2, 4, 9])
@pytest.mark.parametrize("sample_shape", [(), (1, 2), (5, 4, 4)])
def test_buffer(capacity, chunk_len, sample_shape):
  """Builds a Buffer with the provided `capacity` and insert `capacity * 3`
  samples into the buffer in chunks of shape `(chunk_len,) + sample_shape`.
  We always insert chunks with consecutive integers.

  * `len(buffer)` should increase until we reach capacity.
  * `buffer._idx` should loop between 0 and `capacity - 1`.
  * After every insertion, samples should be in expected range, verifying
    FIFO insertion.
  * Mutating the inserted chunk shouldn't mutate the buffer.
  """
  buf = Buffer(capacity, sample_shape=sample_shape, dtypes=float)

  for i in range(0, capacity*3, chunk_len):
    chunk = _fill_chunk(i, chunk_len, sample_shape)
    assert len(buf) == min(i, capacity)
    assert buf._idx == i % capacity
    buf.store(chunk)
    samples = buf.sample(100)
    assert samples.shape == (100,) + sample_shape

    _check_bound(i + chunk_len, capacity, samples)

  # Confirm that buffer is not mutable from inserted sample.
  chunk[:] = np.nan
  assert not np.any(np.isnan(buf._buffer))


@pytest.mark.parametrize("capacity", [30, 60])
@pytest.mark.parametrize("chunk_len", [1, 4, 9])
@pytest.mark.parametrize("obs_shape", [(), (1, 2)])
@pytest.mark.parametrize("act_shape", [(), (5, 4, 4)])
@pytest.mark.parametrize("dtype", [np.int, np.float32])
def test_replay_buffer(capacity, chunk_len, obs_shape, act_shape, dtype):
  """Builds a ReplayBuffer with the provided `capacity` and inserts
  `capacity * 3` observation-action-observation samples into the buffer in
  chunks of length `chunk_len`.

  All chunks are of the appopriate observation or action shape, and contain
  the value fill_val.

  `len(buffer)` should increase until we reach capacity.
  `buffer._idx` should loop between 0 and `capacity - 1`.
  After every insertion, samples should only contain 66.6.
  """
  buf = ReplayBuffer(capacity, obs_shape=obs_shape, act_shape=act_shape,
                     obs_dtype=dtype, act_dtype=dtype)

  for i in range(0, capacity*3, chunk_len):
    assert len(buf) == min(i, capacity)
    for b in [buf._old_obs_buffer, buf._act_buffer, buf._new_obs_buffer]:
      assert b._idx == i % capacity

    old_obs_data = _fill_chunk(i, chunk_len, obs_shape, dtype=dtype)
    new_obs_data = _fill_chunk(3 * capacity + i, chunk_len, obs_shape,
                               dtype=dtype)
    act_data = _fill_chunk(6 * capacity + i, chunk_len, act_shape, dtype=dtype)

    buf.store(old_obs_data, act_data, new_obs_data)

    # Are samples right shape?
    old_obs, acts, new_obs = buf.sample(100)
    assert old_obs.shape == new_obs.shape == (100,) + obs_shape
    assert acts.shape == (100,) + act_shape

    # Are samples right data type?
    assert old_obs.dtype == dtype
    assert acts.dtype == dtype
    assert new_obs.dtype == dtype

    # Are samples in range?
    _check_bound(i + chunk_len, capacity, old_obs)
    _check_bound(i + chunk_len, capacity, new_obs, 3 * capacity)
    _check_bound(i + chunk_len, capacity, acts, 6 * capacity)

    # Are samples in-order?
    old_obs_fill = _get_fill_from_chunk(old_obs)
    new_obs_fill = _get_fill_from_chunk(new_obs)
    act_fill = _get_fill_from_chunk(acts)

    assert np.all(new_obs_fill - old_obs_fill == 3 * capacity), "out of order"
    assert np.all(act_fill - new_obs_fill == 3 * capacity), "out of order"


@pytest.mark.parametrize("sample_shape", [(), (1,), (5, 2)])
def test_buffer_store_errors(sample_shape):
  capacity = 11
  dtype = "float32"

  def buf():
      return Buffer(capacity, sample_shape, dtype)

  # `data` is empty.
  b = buf()
  with pytest.raises(ValueError):
      b.store(np.empty((0,) + sample_shape, dtype=dtype))

  # `data` has too many samples.
  b = buf()
  with pytest.raises(ValueError):
      b.store(np.empty((capacity + 1,) + sample_shape, dtype=dtype))

  # `data` has the wrong sample shape.
  b = buf()
  with pytest.raises(ValueError):
      b.store(np.empty((1, 3, 3, 3, 3), dtype=dtype))


def test_buffer_sample_errors():
  b = Buffer(10, (2, 1), dtypes=bool)
  with pytest.raises(ValueError):
      b.sample(5)

{'fraction_non_alphanumeric': 0.13075567305573915, 'fraction_numerical': 0.020709407358449, 'mean_word_length': 4.049499443826474, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7280136', 'n_tokens_mistral': 3296, 'n_tokens_neox': 3159, 'n_words': 852}