def add_TreeGate_config_args(parser):
    """TreeGate arguments"""

    group = parser.add_argument_group('model_treegate', 'TreeGate model configuration')
    group.add_argument('--depth', type=int, help='Depth of the TreeGate Network.')
    group.add_argument('--dim_reduce_factor', type=int, help='Dimension reduce factor of the branching policy network.')
    group.add_argument('--infimum', type=int, help='Infimum parameter of the branching policy network.')
    group.add_argument('--norm', type=str, default='none', help='Normalization type of the branching policy network.')
    
    return parser