def train(model: ICMNetwork, optimizer: optim.Optimizer, states: Tensor, actions: Tensor, next_states: Tensor, total_train_steps: int = 1000) -> float:
    """
    **Overview**:
        Training function for the ICM model.
    """
    for _ in range(total_train_steps):
        # Perform a forward pass through the network.
        real_next_state_feature, pred_next_state_feature, pred_action = model(states, next_states, actions)

        # Compute the forward loss as the MSE loss between the real and predicted next state features.
        forward_loss = ((real_next_state_feature - pred_next_state_feature) ** 2).mean()
        # Compute the inverse loss as the MSE loss between the real and predicted actions.
        inverse_loss = ((actions - pred_action) ** 2).mean()

        # Combine the forward and inverse losses into a total loss.
        loss = forward_loss + inverse_loss
        # Zero out the gradients before the backward pass.
        optimizer.zero_grad()
        # Perform a backward pass through the network.
        loss.backward()
        # Update the network parameters with a step of the optimizer.
        optimizer.step()

    # Return the final loss.
    return loss.item()

    **Overview**:
        Training function for the ICM model.
    
chapter4_reward/curiosity_icm.py