def METHOD_NAME(self, observation, state, step_type):
    """Update the goal if the episode just beginned; otherwise keep using
    the goal in ``state``.
    Args:
        observation (nested Tensor): the observation at the current time step
        state (nested Tensor): state of this goal generator
        step_type (StepTyp):
    Returns:
        Tensor: a batch of one-hot tensors representing the updated goals.
    """
    new_goal_mask = torch.unsqueeze((step_type == StepType.FIRST), dim=-1)
    generated_goal = self._generate_goal(observation, state)
    new_goal = torch.where(new_goal_mask, generated_goal, state.goal)
    return new_goal