def assemble_state_vec(arm_concat: tf.Tensor, arm_format: str,
                       base_concat=None, base_format=None) -> tf.Tensor:
    """
    Assemble the state/action vector from the arm and base.
    """
    state_vec = tf.zeros(STATE_VEC_LEN, dtype=tf.float32)
    mask_vec = tf.zeros(STATE_VEC_LEN, dtype=tf.float32)

    # Assemble the arm state
    arm_concat = tf.cast(arm_concat, tf.float32)
    arm_format = arm_format.split(',')
    # Use the scatter_nd to avoid the duplicate indices
    state_vec = tf.tensor_scatter_nd_update(
        state_vec, 
        [[STATE_VEC_IDX_MAPPING[name]] for name in arm_format],
        arm_concat
    )
    mask_vec = tf.tensor_scatter_nd_update(
        mask_vec, 
        [[STATE_VEC_IDX_MAPPING[name]] for name in arm_format],
        tf.ones(len(arm_format), dtype=tf.float32)
    )

    # Assemble the base state if exists
    if base_concat is not None:
        base_concat = tf.cast(base_concat, tf.float32)
        base_format = base_format.split(',')
        state_vec = tf.tensor_scatter_nd_update(
            state_vec, 
            [[STATE_VEC_IDX_MAPPING[name]] for name in base_format],
            base_concat
        )
        mask_vec = tf.tensor_scatter_nd_update(
            mask_vec, 
            [[STATE_VEC_IDX_MAPPING[name]] for name in base_format],
            tf.ones(len(base_format), dtype=tf.float32)
        )
    return state_vec, mask_vec

    Assemble the state/action vector from the arm and base.
    
policy/RDT/data/preprocess.py