def test_adder(self, n_step, additional_discount, first, steps,
expected_transitions):
# Create a fake client to record our writes and use it in the adder.
client = test_utils.FakeClient()
adder = adders.NStepTransitionAdder(client, n_step, additional_discount)
# Add all the data up to the final step.
adder.add_first(first)
for step in steps[:-1]:
adder.add(*step)
# Make sure the writer has been created but not closed.
self.assertLen(client.writers, 1)
self.assertFalse(client.writers[0].closed)
# Add the final step.
adder.add(*steps[-1])
# Ending the episode should close the writer. No new writer should yet have
# been created as it is constructed lazily.
self.assertLen(client.writers, 1)
self.assertTrue(client.writers[0].closed)
# Make sure our expected and observed transitions match.
observed_transitions = list(p[1][0] for p in client.writers[0].priorities)
for exp, obs in zip(expected_transitions, observed_transitions):
tree.map_structure(np.testing.assert_array_almost_equal, exp, obs)
# Add the start of a second trajectory.
adder.add_first(first)
adder.add(*steps[0])
# Make sure this creates an open writer.
self.assertLen(client.writers, 2)
self.assertFalse(client.writers[1].closed)
def test_delta_encoded_and_chunk_length(self):
max_sequence_length = 5
client = test_utils.FakeClient()
adder = adders.EpisodeAdder(client,
max_sequence_length,
delta_encoded=True,
chunk_length=max_sequence_length)
# Add an episode.
first, steps = test_utils.make_trajectory(range(max_sequence_length))
adder.add_first(first)
for action, step in steps:
adder.add(action, step)
self.assertTrue(client.writers[0].delta_encoded)
self.assertEqual(max_sequence_length, client.writers[0].chunk_length)
def test_max_sequence_length(self, max_sequence_length):
client = test_utils.FakeClient()
adder = adders.EpisodeAdder(client, max_sequence_length)
first, steps = test_utils.make_trajectory(range(max_sequence_length + 1))
adder.add_first(first)
for action, step in steps[:-1]:
adder.add(action, step)
# We should have max_sequence_length-1 timesteps that have been written,
# where the -1 is due to the dangling observation (ie we have actually
# seen max_sequence_length observations).
self.assertLen(client.writers[0].timesteps, max_sequence_length - 1)
# Adding one more step should raise an error.
with self.assertRaises(ValueError):
action, step = steps[-1]
adder.add(action, step)
# Since the last insert failed it should not affect the internal state.
self.assertLen(client.writers[0].timesteps, max_sequence_length - 1)
def test_adder(self, max_sequence_length):
client = test_utils.FakeClient()
adder = adders.EpisodeAdder(client, max_sequence_length)
# Create a simple trajectory to add.
observations = range(max_sequence_length)
first, steps = test_utils.make_trajectory(observations)
# Add everything up to the final transition.
adder.add_first(first)
for action, step in steps[:-1]:
adder.add(action, step)
if max_sequence_length == 2:
# Nothing has been written since we only have an initial step and a
# final step (corresponding to a single transition).
self.assertEmpty(client.writers)
else:
# No priorities should have been written so far but all timesteps (all
# but the last one) should have been sent to the writer.
self.assertEmpty(client.writers[0].priorities)
self.assertLen(client.writers[0].timesteps, max_sequence_length - 2)
# Adding a terminal timestep should close the writer and insert a
# prioritized sample referencing all the timesteps (including the padded
# terminating observation).
action, step = steps[-1]
adder.add(action, step)
# The writer should be closed and should have max_sequence_length timesteps.
self.assertTrue(client.writers[0].closed)
self.assertLen(client.writers[0].timesteps, max_sequence_length)
# Make the sequence of data and the priority table entry we expect.
expected_sequence = test_utils.make_sequence(observations)
expected_entry = (base.DEFAULT_PRIORITY_TABLE, expected_sequence, 1.0)
self.assertEqual(client.writers[0].priorities, [expected_entry])
def test_adder(self, sequence_length: int, period: int, first, steps,
expected_sequences, pad_end_of_episode: bool = True):
client = test_utils.FakeClient()
adder = adders.SequenceAdder(
client,
sequence_length=sequence_length,
period=period,
pad_end_of_episode=pad_end_of_episode)
# Add all the data up to the final step.
adder.add_first(first)
for step in steps[:-1]:
adder.add(*step)
# Make sure the writer has been created but not closed.
self.assertLen(client.writers, 1)
self.assertFalse(client.writers[0].closed)
# Add the final step.
adder.add(*steps[-1])
# Ending the episode should close the writer. No new writer should yet have
# been created as it is constructed lazily.
self.assertLen(client.writers, 1)
self.assertTrue(client.writers[0].closed)
# Make sure our expected and observed transitions match.
observed_sequences = list(p[1] for p in client.writers[0].priorities)
for exp, obs in zip(expected_sequences, observed_sequences):
np.testing.assert_array_equal(exp, obs)
# Add the start of a second trajectory.
adder.add_first(first)
adder.add(*steps[0])
# Make sure this creates an open writer.
self.assertLen(client.writers, 2)
self.assertFalse(client.writers[1].closed)