def test_ExtractState(self): extracted_state = state_util.extract_state(self._batched_states, 1) self._assert_sructures_equal(self._unbatched_states[1], extracted_state)