def test_base_buffer_exception(self): with self.assertRaises(ValueError): # size <= 0 ub_data.BaseBuffer(0, 1) with self.assertRaises(AssertionError): # batch <= 0 ub_data.BaseBuffer(1, 0) buf = ub_data.BaseBuffer(1, 1) self.assertTrue(buf.isnull) with self.assertRaises(AssertionError): # AssertionError: Buffer space not created buf[0] # auto create space buf._set_data(1, indices=0) buf[0]
def test_base_buffer_relative_index(self): capacity = 10 batch = 1 n_samples = 15 # test circular buf = ub_data.BaseBuffer(capacity, batch=batch) for i in range(n_samples): buf.add({'a': ([i], [i + 1])}) head = n_samples % capacity self.assertEqual(head, buf.head) self.assertEqual(head, buf.tail) # test int, slice key data = buf.rel[1] self.assertArrayEqual([head + 1], data['a'][0]) self.assertArrayEqual([head + 2], data['a'][1]) data = buf.rel[-1] self.assertArrayEqual([n_samples - 1], data['a'][0]) self.assertArrayEqual([n_samples], data['a'][1]) data = buf.rel[1:3] exp = np.arange(2).reshape(-1, 1) self.assertArrayEqual(exp + head + 1, data['a'][0]) self.assertArrayEqual(exp + head + 2, data['a'][1]) data = buf.rel[-3:-1] exp = np.arange(2, 0, -1).reshape(-1, 1) self.assertArrayEqual(n_samples - 1 - exp, data['a'][0]) self.assertArrayEqual(n_samples - exp, data['a'][1]) data = buf.rel[-1:1] self.assertEqual((0, 1), data['a'][0].shape) self.assertEqual((0, 1), data['a'][1].shape) # test tuple key data = buf.rel[1, 0] self.assertArrayEqual(head + 1, data['a'][0]) self.assertArrayEqual(head + 2, data['a'][1]) data = buf.rel[-1, 0] self.assertArrayEqual(n_samples - 1, data['a'][0]) self.assertArrayEqual(n_samples, data['a'][1]) data = buf.rel[1:3, 0] exp = np.arange(2) self.assertArrayEqual(exp + head + 1, data['a'][0]) self.assertArrayEqual(exp + head + 2, data['a'][1]) data = buf.rel[-3:-1, 0] exp = np.arange(2, 0, -1) self.assertArrayEqual(n_samples - 1 - exp, data['a'][0]) self.assertArrayEqual(n_samples - exp, data['a'][1]) data = buf.rel[-1:1, 0] self.assertEqual((0, ), data['a'][0].shape) self.assertEqual((0, ), data['a'][1].shape) # test list key data = buf.rel[[1, 3]] self.assertArrayEqual([[head + 1], [head + 3]], data['a'][0]) self.assertArrayEqual([[head + 2], [head + 4]], data['a'][1]) # test np key data = buf.rel[np.asarray([1, 3])] self.assertArrayEqual([[head + 1], [head + 3]], data['a'][0]) self.assertArrayEqual([[head + 2], [head + 4]], data['a'][1]) # test index out of range with self.assertRaises(IndexError): buf.rel[capacity + 1]
def test_sampler_exception(self): with self.assertRaises(ValueError): # `buffer` must be an instance of BaseBuffer ub_data.UniformSampler(None) with self.assertRaises(ValueError): # `buffer` must be an instance of BaseBuffer ub_data.PermuteSampler(None) with self.assertRaises(ValueError): # `buffer` must be ReplayBuffer ub_data.PrioritizedSampler(ub_data.BaseBuffer(10), 1.0) with self.assertRaises(ValueError): # `buffer` must be ReplayBuffer ub_data.PrioritizedSampler(ub_data.DynamicBuffer(), 1.0)
def test_base_buffer_auto_calc_space(self): capacity = 10 batch = 1 buf = ub_data.BaseBuffer(capacity, batch=batch) self.assertEqual(0, len(buf)) self.assertEqual(0, buf.len_slots()) self.assertEqual(capacity, buf.capacity) self.assertEqual(capacity, buf.slots) self.assertEqual(batch, buf.batch) self.assertEqual(0, buf.head) self.assertEqual(0, buf.tail) self.assertTrue(buf.isnull) self.assertFalse(buf.isfull) self.assertTrue(buf.ready_for_sample) capacity = 10 n_samples = 15 # test circular buf = ub_data.BaseBuffer(capacity, batch=None) self.assertEqual(0, len(buf)) self.assertEqual(0, buf.len_slots()) self.assertEqual(None, buf.capacity) self.assertEqual(None, buf.slots) self.assertEqual(None, buf.batch) self.assertEqual(0, buf.head) self.assertEqual(0, buf.tail) self.assertTrue(buf.isnull) self.assertFalse(buf.isfull) self.assertTrue(buf.ready_for_sample) buf.add({'a': [0, 1]}) self.assertEqual(2, len(buf)) self.assertEqual(1, buf.len_slots()) self.assertEqual(capacity, buf.capacity) self.assertEqual(math.ceil(capacity / 2), buf.slots) self.assertEqual(2, buf.batch) self.assertEqual(0, buf.head) self.assertEqual(1, buf.tail) self.assertFalse(buf.isnull) self.assertFalse(buf.isfull) self.assertTrue(buf.ready_for_sample)
def test_base_buffer_shape(self): capacity = 10 batch = 3 n_samples = 15 buf = ub_data.BaseBuffer(capacity, batch=batch) for i in range(n_samples): buf.add({'a': np.arange(batch)}) self.assertArrayEqual((3, ), buf[1]['a'].shape) self.assertArrayEqual((2, 3), buf[1:3]['a'].shape) self.assertArrayEqual((2, ), buf[1:3, 0]['a'].shape) self.assertArrayEqual((2, 2), buf[1:3, :2]['a'].shape) self.assertArrayEqual((3, ), buf.rel[1]['a'].shape) self.assertArrayEqual((2, 3), buf.rel[1:3]['a'].shape) self.assertArrayEqual((2, ), buf.rel[1:3, 0]['a'].shape) self.assertArrayEqual((2, 2), buf.rel[1:3, :2]['a'].shape)
def test_uniform_sampler_with_base_buffer(self): capacity = 10 batch = 1 n_samples = 15 buf = ub_data.BaseBuffer(capacity, batch=batch) samp = ub_data.UniformSampler(buf) for i in range(n_samples): buf.add({'a': ([i], [i + 1])}) if i < capacity - 1: self.assertFalse(buf.isfull) self.assertEqual(i + 1, len(buf)) else: self.assertTrue(buf.isfull) self.assertEqual(capacity, len(buf)) exp = np.arange(n_samples - capacity, n_samples) exp_a0 = np.roll(exp, n_samples % capacity) exp_a1 = exp_a0 + 1 exp_a0 = np.expand_dims(exp_a0, axis=-1) exp_a1 = np.expand_dims(exp_a1, axis=-1) self.assertArrayEqual(buf.data['a'][0], exp_a0) self.assertArrayEqual(buf.data['a'][1], exp_a1) # test sample (batch=None) batch = samp() self.assertArrayEqual((capacity, ), batch['a'][0].shape) self.assertArrayEqual((capacity, ), batch['a'][1].shape) self.assertArrayEqual(batch['a'][0], buf[samp.indices]['a'][0]) self.assertArrayEqual(batch['a'][1], buf[samp.indices]['a'][1]) # test sample (batch=3) batch_size = 3 batch = samp(batch_size=batch_size) self.assertArrayEqual((batch_size, ), batch['a'][0].shape) self.assertArrayEqual((batch_size, ), batch['a'][1].shape) self.assertArrayEqual(batch['a'][0], buf[samp.indices]['a'][0]) self.assertArrayEqual(batch['a'][1], buf[samp.indices]['a'][1]) # test sample (batch=3, seq=2) batch_size = 3 seq_len = 2 batch = samp(batch_size=batch_size, seq_len=seq_len) self.assertArrayEqual((batch_size, seq_len), batch['a'][0].shape) self.assertArrayEqual((batch_size, seq_len), batch['a'][1].shape) self.assertArrayEqual(batch['a'][0], buf[samp.indices]['a'][0]) self.assertArrayEqual(batch['a'][1], buf[samp.indices]['a'][1]) # test update batch['a'] = (np.zeros_like(batch['a'][0]), np.zeros_like(batch['a'][1])) samp.update(batch) self.assertArrayEqual(buf[samp.indices]['a'][0], batch['a'][0]) self.assertArrayEqual(buf[samp.indices]['a'][1], batch['a'][1])
def test_uniform_sampler_with_base_buffer_rel(self): capacity = 10 batch = 1 n_samples = 15 buf = ub_data.BaseBuffer(capacity, batch=batch) samp = ub_data.UniformSampler(buf) for i in range(n_samples): buf.add({'a': ([i], [i + 1])}) # test relative indexing inds1 = np.arange(3) inds2 = np.zeros(3, dtype=np.int64) samp._cached_inds = (inds1, inds2) self.assertArrayEqual(buf[samp.indices]['a'][0], samp.rel[0]['a'][0]) self.assertArrayEqual(buf[samp.indices]['a'][1], samp.rel[0]['a'][1]) self.assertArrayEqual(buf[(inds1 - 3, inds2)]['a'][0], samp.rel[-3]['a'][0]) self.assertArrayEqual(buf[(inds1 - 3, inds2)]['a'][1], samp.rel[-3]['a'][1]) self.assertArrayEqual(buf[(inds1 + 3, inds2)]['a'][0], samp.rel[3]['a'][0]) self.assertArrayEqual(buf[(inds1 + 3, inds2)]['a'][1], samp.rel[3]['a'][1]) add = np.array([[1, 2]], dtype=np.int64).T self.assertArrayEqual(buf[(inds1 + add, inds2)]['a'][0], samp.rel[1:3]['a'][0]) self.assertArrayEqual(buf[(inds1 + add, inds2)]['a'][1], samp.rel[1:3]['a'][1]) add = np.array([[-3, -2]], dtype=np.int64).T self.assertArrayEqual(buf[(inds1 + add, inds2)]['a'][0], samp.rel[-3:-1]['a'][0]) self.assertArrayEqual(buf[(inds1 + add, inds2)]['a'][1], samp.rel[-3:-1]['a'][1]) # test shape self.assertArrayEqual((0, 3), samp.rel[-1:1]['a'][0].shape) # test setitem add = np.array([[1, 2]], dtype=np.int64).T batch = samp.rel[1:3] batch['a'] = (np.zeros_like(batch['a'][0]), np.zeros_like(batch['a'][1])) samp.rel[1:3] = batch self.assertTrue(np.all(samp.rel[1:3]['a'][0] == 0)) self.assertTrue(np.all(samp.rel[1:3]['a'][1] == 0))
def test_base_buffer_multidim(self): capacity = 20 batch = 2 dim = 2 n_samples = 15 # test circular buf = ub_data.BaseBuffer(capacity, batch=batch) data = np.arange(n_samples * batch * dim).reshape( (n_samples, batch, dim)) for i in range(n_samples): buf.add({'a': data[i]}) if (i + 1) * batch < capacity: self.assertFalse(buf.isfull) self.assertEqual((i + 1) * batch, len(buf)) self.assertEqual(i + 1, buf.len_slots()) self.assertEqual(0, buf.head) else: self.assertTrue(buf.isfull) self.assertEqual(capacity, len(buf)) self.assertEqual(capacity // batch, buf.len_slots()) self.assertEqual((i + 1) % (capacity // batch), buf.tail) exp = np.arange(n_samples * batch * dim - capacity * dim, n_samples * batch * dim) exp = exp.reshape(-1, 2, 2) exp = np.roll(exp, n_samples % (capacity // batch), axis=0) self.assertArrayEqual(exp, buf.data['a']) # test ravel/unravel index def test_ravel(indices): self.assertArrayEqual( np.ravel_multi_index(indices, (buf.slots, buf.batch)), buf.ravel_index(indices)) test_ravel(([1, 2, 3], 0)) test_ravel(([[1], [2], [3]], [0, 1])) def test_unravel(indices): self.assertArrayEqual( np.unravel_index(indices, (buf.slots, buf.batch)), buf.unravel_index(indices)) test_unravel([4, 5, 6]) test_unravel(7)
def test_permute_sampler_with_base_buffer(self): capacity = 10 batch = 1 n_samples = 15 buf = ub_data.BaseBuffer(capacity, batch=batch) samp = ub_data.PermuteSampler(buf) for i in range(n_samples): buf.add({'a': ([i], [i + 1])}) # test sample (batch=None) batch_size = len(buf) batches = [] indices = [] for batch in samp(): self.assertArrayEqual((batch_size, ), batch['a'][0].shape) self.assertArrayEqual((batch_size, ), batch['a'][1].shape) self.assertArrayEqual(buf[samp.indices]['a'][0], batch['a'][0]) self.assertArrayEqual(buf[samp.indices]['a'][1], batch['a'][1]) batches.append(batch) indices.append( np.ravel_multi_index(samp.indices, (buf.len_slots(), buf.batch))) self.assertEqual(1, len(batches)) unique, counts = np.unique(indices, return_counts=True) # check if contains all elements self.assertTrue(len(buf), len(unique)) # check if all elements are sampled at least once self.assertTrue(np.all(counts == 1)) # test sample (batch=3) batch_size = 3 batches = [] indices = [] for batch in samp(batch_size=batch_size): self.assertArrayEqual((batch_size, ), batch['a'][0].shape) self.assertArrayEqual((batch_size, ), batch['a'][1].shape) self.assertArrayEqual(buf[samp.indices]['a'][0], batch['a'][0]) self.assertArrayEqual(buf[samp.indices]['a'][1], batch['a'][1]) batches.append(batch) indices.append( np.ravel_multi_index(samp.indices, (buf.len_slots(), buf.batch))) self.assertEqual(4, len(batches)) # total samples == capacity unique, counts = np.unique(indices, return_counts=True) # check if contains all elements self.assertTrue(len(buf), len(unique)) # check if all elements are sampled at least once but less than 2 self.assertTrue(np.all(counts >= 1)) self.assertTrue(np.all(counts <= 2)) # test sample (batch=3, seq_len=2) batch_size = 3 seq_len = 2 batches = [] indices = [] for batch in samp(batch_size=batch_size, seq_len=seq_len): self.assertArrayEqual((batch_size, seq_len), batch['a'][0].shape) self.assertArrayEqual((batch_size, seq_len), batch['a'][1].shape) self.assertArrayEqual(buf[samp.indices]['a'][0], batch['a'][0]) self.assertArrayEqual(buf[samp.indices]['a'][1], batch['a'][1]) batches.append(batch) indices.append( np.ravel_multi_index(samp.indices, (buf.len_slots(), buf.batch))) self.assertEqual(3, len(batches)) # total samples == capacity unique, counts = np.unique(indices, return_counts=True) # check if contains all elements self.assertTrue(len(buf), len(unique)) # check if all elements are sampled at least once but less than 3 self.assertTrue(np.all(counts >= 1)) self.assertTrue(np.all(counts <= 4))
def test_base_buffer(self): capacity = 10 batch = 1 n_samples = 15 # test circular buf = ub_data.BaseBuffer(capacity, batch=batch) self.assertEqual(capacity, buf.capacity) self.assertEqual(capacity, buf.slots) self.assertEqual(batch, buf.batch) self.assertEqual(0, buf.head) self.assertEqual(0, buf.tail) self.assertTrue(buf.isnull) self.assertFalse(buf.isfull) self.assertTrue(buf.ready_for_sample) for i in range(n_samples): buf.add({'a': ([i], [i + 1])}) if i < capacity - 1: self.assertFalse(buf.isfull) self.assertEqual(i + 1, len(buf)) self.assertEqual(i + 1, buf.len_slots()) self.assertEqual(0, buf.head) else: self.assertTrue(buf.isfull) self.assertEqual(capacity, len(buf)) self.assertEqual(capacity, buf.len_slots()) self.assertEqual((i + 1) % capacity, buf.tail) exp = np.arange(n_samples - capacity, n_samples) exp_a0 = np.roll(exp, n_samples % capacity) exp_a1 = exp_a0 + 1 exp_a0 = np.expand_dims(exp_a0, axis=-1) exp_a1 = np.expand_dims(exp_a1, axis=-1) self.assertArrayEqual(exp_a0, buf.data['a'][0]) self.assertArrayEqual(exp_a1, buf.data['a'][1]) # test getitem data = buf[np.arange(n_samples % capacity)] exp_a0 = np.arange(n_samples - n_samples % capacity, n_samples) exp_a1 = exp_a0 + 1 exp_a0 = np.expand_dims(exp_a0, axis=-1) exp_a1 = np.expand_dims(exp_a1, axis=-1) self.assertArrayEqual(exp_a0, data['a'][0]) self.assertArrayEqual(exp_a1, data['a'][1]) # test setitem n = n_samples - capacity new_data = np.arange(n - n_samples % capacity, n) new_data = np.expand_dims(new_data, axis=-1) new_data = {'a': (new_data, new_data + 1)} buf[np.arange(n_samples % capacity)] = new_data n = n_samples - capacity - n_samples % capacity exp_a0 = np.arange(n, n + capacity) exp_a1 = exp_a0 + 1 exp_a0 = np.expand_dims(exp_a0, axis=-1) exp_a1 = np.expand_dims(exp_a1, axis=-1) self.assertArrayEqual(exp_a0, buf.data['a'][0]) self.assertArrayEqual(exp_a1, buf.data['a'][1]) # test update (should have the same results as setitem) buf.update(new_data, indices=np.arange(n_samples % capacity)) self.assertArrayEqual(exp_a0, buf.data['a'][0]) self.assertArrayEqual(exp_a1, buf.data['a'][1]) # test ravel/unravel index def test_ravel(indices): self.assertArrayEqual( np.ravel_multi_index(indices, (buf.slots, buf.batch)), buf.ravel_index(indices)) test_ravel(([1, 2, 3], 0)) test_ravel(([1, 2, 3], [0])) def test_unravel(indices): self.assertArrayEqual( np.unravel_index(indices, (buf.slots, buf.batch)), buf.unravel_index(indices)) test_unravel([4, 5, 6]) test_unravel(7)