A.- B._ a bénéficié d'une demi-rente d'invalidité pour son fils A._. Par décision du 31 juillet 2000, l'Office de l'assurance-invalidité pour le canton de Vaud a supprimé cette prestation, avec effet dès ce jour-là, au motif que A._ avait terminé ses études. B.- B._ a recouru contre cette décision devant le Tribunal des assurances du canton de Vaud. A l'appui de son recours, elle a produit une attestation de l'Association Y._ du 17 août 2000, dont il ressort que son fils participait à un semestre de motivation du 14 août 2000 au 13 février 2001. Par jugement du 21 décembre 2000, la juridiction cantonale a rejeté le recours. C.- A._ et B._ interjettent recours de droit administratif contre ce jugement dont ils demandent l'annulation. L'office intimé conclut au rejet du recours, en se référant au préavis de la Caisse cantonale vaudoise de compensation (la caisse). L'Office fédéral des assurances sociales (OFAS) propose de rejeter le recours, au terme d'un préavis sur lequel les parties ont eu l'occasion de se déterminer. 