def test_sample_multi_discrete_action():
    """
    **Overview**:
        The function of testing sampling multi-discrete action. Construct a naive policy and sample a group of multi-discrete action.
    """
    # Set batch_size = 4, obs_shape = 10, action_shape = [4, 5, 6].
    B, obs_shape, action_shape = 4, 10, [4, 5, 6]
    # Generate state data from uniform distribution in [0, 1].
    state = torch.rand(B, obs_shape)
    # Define policy network with encoder and head.
    policy_network = MultiDiscretePolicyNetwork(obs_shape, action_shape)
    # Policy network forward procedure, input state and output multiple logit.
    # $$ logit = \pi(a|s)$$
    logit = policy_network(state)
    for i in range(len(logit)):
        assert logit[i].shape == (B, action_shape[i])
    # Sample action accoding to corresponding logit one by one.
    for i in range(len(logit)):
        action_i = sample_action(logit[i])
        assert action_i.shape == (B, )

    **Overview**:
        The function of testing sampling multi-discrete action. Construct a naive policy and sample a group of multi-discrete action.
    
chapter2_action/discrete_tutorial.py