def extend_model_config(
    model_config, syndrome_size, stack_depth, num_actions_per_qubit=3, device="cpu"
):
    """
    Extend an existing model or agent configuration dictionary
    with information about the environment.

    Parameters
    ==========
    model_config: (dict) dictionary contiaining information about
        model architecture and layer shapes
    syndrome_size: (int) size of the state, ususally code distance+1
    stack_depth: (int) number of layers in a state stack
    num_actions_per_qubit: (optional) (int), number of possible actions on one
        qubit. Defaults to 3 for Pauli-X, -Y, -Z.

    Returns
    =======
    model_config: (dict) updated dictionary with configuration information
        of the model architecture
    """

    model_config["syndrome_size"] = syndrome_size
    model_config["code_size"] = syndrome_size - 1
    model_config["stack_depth"] = stack_depth
    model_config["num_actions_per_qubit"] = num_actions_per_qubit
    model_config["device"] = device

    return model_config