def forward(self, local_obs: torch.Tensor) -> ttorch.Tensor:
        """
        **Overview**:
            The computation graph of actor-critic network in discrete action space.
        """
        # Transform original local obs into embedding vector, i.e. $$(B, A, *) -> (B, A, N)$$
        # Some network layers in PyTorch like ``nn.Linear`` can deal with any number of prefix dimensions, so we can just use it to process the whole multi-agent batch.
        x = self.encoder(local_obs)
        # Calculate logit for each possible discrete action choices, i.e. $$(B, A, N) -> (B, A, M)$$
        logit = self.policy_head(x)
        # Calculate value for each sample and agent, i.e. $$(B, A, N) -> (B, A, 1)$$
        value = self.value_head(x)
        # Return the final result by treetensor format.
        return ttorch.as_tensor({
            'logit': logit,
            'value': value,
        })

        **Overview**:
            The computation graph of actor-critic network in discrete action space.
        
chapter6_marl/marl_network.py