def add_training_args(parser):
    """Training arguments."""

    group = parser.add_argument_group('train', 'training configurations')
    group.add_argument('--train_batchsize', type=int, default=32, help='Training batchsize.')
    group.add_argument('--seed', type=int, default=0, help='Random seed for IL training experiment.')
    group.add_argument('--opt', default='adam', type=str, help='Type of optimizer to use.')
    group.add_argument('--lr', type=float, default=1e-5, help='Learning rate.')
    group.add_argument('--momentum', default=0.9, type=float, help='Momentum optimization parameter.')
    group.add_argument('--weight_decay', default=1e-5, type=float, help='Weight decay optimization parameter.')
    group.add_argument('--num_epochs', type=int, default=40, help='Number of training epochs.')
    group.add_argument('--top_k', type=int, nargs='+', default=[2, 5],
        help='In addition to top-1 generalization accuracy, we track top-k.'
    )
    group.add_argument('--use_gpu', default=False, action='store_true', help='Use gpu or not.')
    group.add_argument('--out_dir', type=str, default='./output', help='Directory to save the experimental results.')

    group.add_argument('--hidden_size', type=int, help='Hidden size of the branching policy network.')
    group.add_argument('--dropout', type=float, default=0.0, help='Dropout parameter for the branching policy network.')

    # TreeGate
    group.add_argument('--lr_decay_schedule', type=int, nargs='+', default=[20, 30],
        help='Learning rate decay schedule.')
    group.add_argument('--lr_decay_factor', type=float, default=0.1, help='LR decay factor.')

    # T-BranT
    group.add_argument('--noam', default=False, action='store_true', help='Use the Noam Scheduler.')
    group.add_argument('--warm_epochs', type=int, default=5, help='Warm epochs for Noam Scheduler.')

    return parser