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