def test_sample_discrete_action():
    """
    **Overview**:
        The function of testing sampling discrete action. Construct a naive policy and sample a group of action.
    """
    # Set batch_size = 4, obs_shape = 10, action_shape = 6.
    B, obs_shape, action_shape = 4, 10, 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 = DiscretePolicyNetwork(obs_shape, action_shape)
    # Policy network forward procedure, input state and output logit.
    # $$ logit = \pi(a|s)$$
    logit = policy_network(state)
    assert logit.shape == (B, action_shape)
    # Sample action accoding to corresponding logit.
    action = sample_action(logit)
    assert action.shape == (B, )

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