def reset_with_initial_state(env, index):
      """Resets environment taking self._initial_state into account."""
      obs = env.reset()
      if self._initial_state is None:
        return obs
      else:
        env.unwrapped.restore_full_state(self._initial_state[index])
        return self._initial_frames[index, ...]
Resets environment taking self._initial_state into account.
tensor2tensor/data_generators/gym_env.py