def train(obs_shape: int, action_shape: int, NUM_EPOCH: int, train_data):
    """
    **Overview**:
        Example training function for using MLP network with Pop-Art layer in fixed Q value approximation.
    """
    # Define the MLP network and optimizer, and loss function.
    model = MLP(obs_shape, action_shape)
    optimizer = AdamW(model.parameters(), lr=0.0001, weight_decay=0.0001)
    MSEloss = nn.MSELoss()
    # Read the preprocessed data of trained agent on lunarlander environment.
    # Each sample in the datasets should be a dict with following format:
    # $$key\quad dim$$
    # $$observations\quad (*,8)$$
    # $$actions\quad (*,)$$
    # $$returns\quad (*,)$$
    # where the returns is the discounted return from the current state.
    train_data = DataLoader(train_data, batch_size=64, shuffle=True)

    # For loop 1: train MLP network for ``NUM_EPOCH`` epochs.
    running_loss = 0.0
    for epoch in range(NUM_EPOCH):
        # For loop 2: Inside each epoch, split the entire dataset into mini-batches, then train on each mini-batch.
        for idx, data in enumerate(train_data):
            # Compute the original output and the normalized output.
            output = model(data['observations'], data['actions'])
            mu = model.popart.mu
            sigma = model.popart.sigma
            # Normalize the target return to align with the normalized Q value.
            with torch.no_grad():
                normalized_return = (data['returns'] - mu) / sigma
            # The loss is calculated as the MSE loss between normalized Q value and normalized target return.
            loss = MSEloss(output.normalized_output, normalized_return)
            # Loss backward and optimizer update step.
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # After the model parameters are updated with the gradient,
            # the weights and bias should be updated to preserve unnormalised outputs.
            model.popart.update_parameters(data['returns'])

            # Use ``item`` method to get the pure Python scalar of the loss, then add it into ``running_loss`` .
            running_loss += loss.item()

        # Print the loss every 100 epochs.
        if epoch % 100 == 99:
            print('Epoch [%d] loss: %.6f' % (epoch + 1, running_loss / 100))
            running_loss = 0.0

    **Overview**:
        Example training function for using MLP network with Pop-Art layer in fixed Q value approximation.
    
chapter4_reward/popart.py