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.env.set_initial_state(initial_state, initial_frames)
    self._initial_frames = initial_frames