Sets the state that will be used on next reset.

def set_initial_state(self, initial_state, initial_frames):
    """Sets the state that will be used on next reset."""
    self._initial_state = initial_state
    self._initial_frames = initial_frames[:, -1, ...]
    self._should_preprocess_on_reset = False