def sample_action(logit: torch.Tensor) -> torch.Tensor:
    """
    **Overview**:
        The function of sampling discrete action, input shape = (B, action_shape), output shape = (B, ).
        In this example, the distributions shapes are:
        batch_shape = (B, ), event_shape = (), sample_shape = ().
    """
    # Transform logit (raw output of policy network, e.g. last fully connected layer) into probability.
    # $$\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$$
    prob = torch.softmax(logit, dim=-1)
    # Construct categorical distribution. The probability mass function is: $$f(x=i|\boldsymbol{p})=p_i$$
    # <link https://en.wikipedia.org/wiki/Categorical_distribution link>
    dist = torch.distributions.Categorical(probs=prob)
    # Sample one discrete action per sample (state input) and return it.
    return dist.sample()

    **Overview**:
        The function of sampling discrete action, input shape = (B, action_shape), output shape = (B, ).
        In this example, the distributions shapes are:
        batch_shape = (B, ), event_shape = (), sample_shape = ().
    
chapter2_action/discrete_tutorial.py