def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        **Overview**:
            The computation graph of hybrid action policy network used in PPO.
            ``x -> encoder -> action_type_head -> action_type_logit`` .
            ``x -> encoder -> action_args_mu -> \mu`` .
            ``action_args_log_sigma -> exp -> sigma`` .
        """
        # Transform original state into embedding vector, i.e. $$(B, *) -> (B, N)$$
        x = self.encoder(x)
        # Output discrete action logit.
        logit = self.action_type_head(x)
        # Output the argument mu depending on the embedding vector.
        mu = self.action_args_mu(x)
        # Utilize broadcast mechanism to make the same shape between log_sigma and mu.
        # ``zeros_like`` operation doesn't pass gradient.
        # <link https://pytorch.org/tutorials/beginner/introyt/tensors_deeper_tutorial.html#in-brief-tensor-broadcasting link>
        log_sigma = self.action_args_log_sigma + torch.zeros_like(mu)
        # Utilize exponential operation to produce the actual sigma.
        # $$\sigma = e^w$$
        sigma = torch.exp(log_sigma)
        # Return treetensor-type output.
        return ttorch.as_tensor({'action_type': logit, 'action_args': {'mu': mu, 'sigma': sigma}})

        **Overview**:
            The computation graph of hybrid action policy network used in PPO.
            ``x -> encoder -> action_type_head -> action_type_logit`` .
            ``x -> encoder -> action_args_mu -> \mu`` .
            ``action_args_log_sigma -> exp -> sigma`` .
        
chapter2_action/hybrid_tutorial.py