def sample_hybrid_action(logit: ttorch.Tensor) -> torch.Tensor:
    """
    **Overview**:
        The function of sampling hybrid action, input is a treetensor with two keys ``action_type`` and ``action_args`` .
    """
    # Transform logit (raw output of discrete policy head, 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.action_type, 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>
    discrete_dist = torch.distributions.Categorical(probs=prob)
    # Sample one discrete action type per sample (state input).
    action_type = discrete_dist.sample()

    # Construct gaussian distribution
    # $$X \sim \mathcal{N}(\mu,\,\sigma^{2})$$
    # Its probability density function is: $$f(x) = \frac{1}{\sigma\sqrt{2\pi}} \exp\left( -\frac{1}{2}\left(\frac{x-\mu}{\sigma}\right)^{\!2}\,\right)$$
    # <link https://en.wikipedia.org/wiki/Normal_distribution link>
    continuous_dist = Normal(logit.action_args.mu, logit.action_args.sigma)
    # Reinterpret ``action_shape`` gaussian distribution into a multivariate gaussian distribution with
    # diagonal convariance matrix.
    # Ensure each event is independent with each other.
    # <link https://pytorch.org/docs/stable/distributions.html#independent link>
    continuous_dist = Independent(continuous_dist, 1)
    # Sample one action args of the shape ``action_shape`` per sample (state input).
    action_args = continuous_dist.sample()
    # Return the final parameterized action.
    return ttorch.as_tensor({
        'action_type': action_type,
        'action_args': action_args,
    })

    **Overview**:
        The function of sampling hybrid action, input is a treetensor with two keys ``action_type`` and ``action_args`` .
    
chapter2_action/hybrid_tutorial.py