def test_sample_continuous_action():
    """
    **Overview**:
        The function of testing sampling continuous action. Construct a standard continuous action
        policy and sample a group of action.
    """
    # Set batch_size = 4, obs_shape = 10, action_shape = 6.
    # ``action_shape`` is different from discrete and continuous action. The former is the possible
    # choice of a discrete action while the latter is the dimension of continuous action.
    B, obs_shape, action_shape = 4, 10, 6
    # Generate state data from uniform distribution in [0, 1].
    state = torch.rand(B, obs_shape)
    # Define continuous action network (which is similar to reparameterization) with encoder, mu and log_sigma.
    policy_network = ContinuousPolicyNetwork(obs_shape, action_shape)
    # Policy network forward procedure, input state and output dict-type logit.
    # $$ \mu, \sigma = \pi(a|s)$$
    logit = policy_network(state)
    assert isinstance(logit, dict)
    assert logit['mu'].shape == (B, action_shape)
    assert logit['sigma'].shape == (B, action_shape)
    # Sample action accoding to corresponding logit (i.e., mu and sigma).
    action = sample_continuous_action(logit)
    assert action.shape == (B, action_shape)

    **Overview**:
        The function of testing sampling continuous action. Construct a standard continuous action
        policy and sample a group of action.
    
chapter2_action/continuous_tutorial.py