def testCheckpoint(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) empty_values = self.np_values(spec, 0) test_values = self.np_values(spec, 5) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(replay_table.append(1, test_values)) checkpoint = tf.train.Checkpoint(table=replay_table) tmpdir = self.get_temp_dir() if tf.executing_eagerly(): location = checkpoint.save(os.path.join(tmpdir, 'ckpt')) else: with self.cached_session() as sess: location = checkpoint.save(os.path.join(tmpdir, 'ckpt'), session=sess) reload_replay_table = episodic_table.EpisodicTable(spec, capacity=3) reload_checkpoint = tf.train.Checkpoint(table=reload_replay_table) status = reload_checkpoint.restore(location) status.assert_consumed() with self.cached_session() as sess: status.initialize_or_restore(session=sess) self._test_episode_lists(replay_table, empty_values, test_values) self._test_episode_lists(reload_replay_table, empty_values, test_values)
def testGetAddMultiple(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=4) batch_size = 2 input_values = self.np_values(spec, batch_size) expected_values = self.np_values(spec) empty_values = self.np_values(spec, 0) tensors = tf.nest.map_structure( lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32), input_values) write_op = replay_table.add([0, 1], tensors) values_0 = replay_table.get_episode_values(0) values_1 = replay_table.get_episode_values(1) # This should be empty values_2 = replay_table.get_episode_values(2) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(write_op) tf.nest.map_structure(self.assertAllClose, expected_values, self.evaluate(values_0)) tf.nest.map_structure(self.assertAllClose, expected_values, self.evaluate(values_1)) tf.nest.map_structure(self.assertAllClose, empty_values, self.evaluate(values_2))
def testGetEpisodeListsSomeRows(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) empty_values = self.np_values(spec, 0) test_values = self.np_values(spec, 5) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(replay_table.append(1, test_values)) # Get episode lists for rows 1 and 2. The returned values should look like: # {'action': [<tensorlist for row 1>, <tensorlist for row 2>], # 'camera': [<tensorlist for row 1>, <tensorlist for row 2>], # 'lidar': [<tensorlist for row 1>, <tensorlist for row 2>]} episode_lists = replay_table.get_episode_lists([1, 2]) for episode_list_slot in tf.nest.flatten(episode_lists): self.assertEqual(episode_list_slot.shape.rank, 1) self.assertEqual(self.evaluate(tf.size(input=episode_list_slot)), 2) # Stack episode tensors for row 1, i.e.: # {'action': <five items>, 'camera': <five items>, 'lidar': <five items>} episode_tensors_1 = tf.nest.map_structure( lambda slot, lists: replay_table._stack_tensor_list( slot, lists[0]), replay_table.slots, episode_lists) # Should be equivalent to # Stack episode tensors for row 2, which is empty. episode_tensors_2 = tf.nest.map_structure( lambda slot, lists: replay_table._stack_tensor_list( slot, lists[1]), replay_table.slots, episode_lists) self.assertAllClose(self.evaluate(episode_tensors_1), test_values) self.assertAllClose(self.evaluate(episode_tensors_2), empty_values)
def testGetEmpty(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) empty_values = self.np_values(spec, 0) values = replay_table.get_episode_values(0) self.evaluate(tf.compat.v1.global_variables_initializer()) values_ = self.evaluate(values) tf.nest.map_structure(self.assertAllClose, empty_values, values_)
def testGetEpisodeListsAllRows(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) empty_values = self.np_values(spec, 0) test_values = self.np_values(spec, 5) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(replay_table.append(1, test_values)) # Get episode lists for all rows. The returned values should look like: # {'action': [<tensorlist for row 0>, <tensorlist for row 1>, # <tensorlist for row 2>], # ... } self._test_episode_lists(replay_table, empty_values, test_values)
def testGetEpisodeListsOneRow(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) test_values = self.np_values(spec, 5) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(replay_table.append(1, test_values)) # Get the episode lists we just added. The returned value should look like: # {'action': <tensorlist>, 'camera': <tensorlist>, 'lidar': <tensorlist>} episode_lists = replay_table.get_episode_lists(1) for episode_list_slot in tf.nest.flatten(episode_lists): self.assertEqual(episode_list_slot.shape.rank, 0) episode_tensors = tf.nest.map_structure( replay_table._stack_tensor_list, replay_table.slots, episode_lists) self.assertAllClose(self.evaluate(episode_tensors), test_values)
def testExtend(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) test_values = self.np_values(spec, 5) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(replay_table.append(1, test_values)) # Extend row 2 by the contents of row 1. self.evaluate(replay_table.extend(2, replay_table.get_episode_lists(1))) # Extend rows 0 and 1 by the contents of rows 1 and 2. self.evaluate( replay_table.extend([0, 1], replay_table.get_episode_lists([1, 2]))) episode_0, episode_1, episode_2 = self.evaluate( [replay_table.get_episode_values(r) for r in range(3)]) self.assertAllClose(episode_0, self.np_values(spec, 5)) self.assertAllClose(episode_1, self.np_values(spec, 10)) self.assertAllClose(episode_2, self.np_values(spec, 5))
def testGetAddSingle(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=3) expected_values = self.np_values(spec) tensors = tf.nest.map_structure( lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32), expected_values) add_op = replay_table.add([0], tensors) values = replay_table.get_episode_values(0) # Check static shape assert_same_shape = lambda s, v: self.assertEqual(s.shape[1:], v.shape) tf.nest.map_structure(assert_same_shape, values, spec) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(add_op) values_ = self.evaluate(values) tf.nest.map_structure(self.assertAllClose, expected_values, values_)
def testGetAddAppendMultiple(self): spec = self.default_specs() replay_table = episodic_table.EpisodicTable(spec, capacity=4) batch_size = 2 input_values = self.np_values(spec, batch_size) expected_values = self.np_values(spec) empty_values = self.np_values(spec, 0) tensors = tf.nest.map_structure( lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32), input_values) # Pull out the first entry in the batch, and add an outer # dimension to represent a single time step that we'll append. tensors_batch0 = tf.nest.map_structure( lambda x: tf.expand_dims(x[0, ...], 0), tensors) # We will append tensors_batch0 to row 0, which contains x[0]. expected_appended_values = tf.nest.map_structure( lambda x: np.stack((x[0], x[0])), input_values) # batch_size == 2, so add [0, 1] write_op = replay_table.add([0, 1], tensors) append_op_0 = lambda: replay_table.append(0, tensors_batch0) values_0 = lambda: replay_table.get_episode_values(0) values_1 = lambda: replay_table.get_episode_values(1) values_2 = lambda: replay_table.get_episode_values(2) self.evaluate(tf.compat.v1.global_variables_initializer()) self.evaluate(write_op) tf.nest.map_structure(self.assertAllClose, expected_values, self.evaluate(values_0())) tf.nest.map_structure(self.assertAllClose, expected_values, self.evaluate(values_1())) tf.nest.map_structure(self.assertAllClose, empty_values, self.evaluate(values_2())) self.evaluate(append_op_0()) tf.nest.map_structure(self.assertAllClose, expected_appended_values, self.evaluate(values_0()))