def train(model: RNDNetwork, optimizer: optim.Optimizer, obs: Tensor, total_train_steps: int = 1000) -> float:
    """
    **Overview**:
        Training function for the RND model.
    """
    # The model is trained for a fixed number of steps
    for _ in range(total_train_steps):
        # Forward pass
        target_feature, pred_feature = model(obs)

        # Compute the loss as the MSE between the target and predicted features
        loss = ((target_feature - pred_feature) ** 2).mean()

        # Backward pass and optimization step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

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

    **Overview**:
        Training function for the RND model.
    
chapter4_reward/curiosity_rnd.py